leetcode 1631. 最小体力消耗路径 python 并查集

你准备参加一场远足活动。给你一个二维 rows x columns 的地图 heights ，其中 heights[row][col] 表示格子 (row, col) 的高度。一开始你在最左上角的格子 (0, 0) ，且你希望去最右下角的格子 (rows-1, columns-1) （注意下标从 0 开始编号）。你每次可以往 上，下，左，右 四个方向之一移动，你想要找到耗费 体力 最小的一条路径。

一条路径耗费的 体力值 是路径上相邻格子之间 高度差绝对值 的 最大值 决定的。

请你返回从左上角走到右下角的最小 体力消耗值 。

class UnionFindSet():
    def __init__(self,n):
        self.setSize = n  #不连通区域
        self.father = {}

    def add(self,x):  #添加根节点
        if x not in self.father:
            self.father[x] = -1
            self.setSize += 1

    def find(self,x):
        root = x
        while self.father[root] != -1:  #找根节点
            root = self.father[root]

        while (x != root): #路径压缩
            o = self.father[x]  #找x的父节点
            self.father[x] = root  #把x的父节点设置成刚才找到的根
            x = o  #往上一层

        return root

    def merge(self, x, y):
        self.add(x)
        self.add(y)
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.father[root_x] = root_y  #合并
            self.setSize -= 1
            
    def is_Union(self, x, y):
        return self.find(x) == self.find(y)

class Solution:
    def minimumEffortPath(self, heights: List[List[int]]) -> int:
        m, n = len(heights), len(heights[0])
        edges = list()
        for i in range(m):
            for j in range(n):
                iden = i * n + j
                if i > 0:
                    edges.append((iden - n, iden, abs(heights[i][j] - heights[i - 1][j])))
                if j > 0:
                    edges.append((iden - 1, iden, abs(heights[i][j] - heights[i][j - 1])))
            
        edges.sort(key=lambda e: e[2])

        uf = UnionFindSet(m * n)
        uf.add(0)
        uf.add(m * n - 1)
        ans = 0
        for x, y, v in edges:
            uf.merge(x, y)
            if uf.is_Union(0, m * n - 1):
                ans = v
                break
            
        return ans