带权重的并查集算法应用

关于并查集算法的应用
并查集算法应用

399. 除法求值

class Solution:
    def calcEquation(self, equations: List[List[str]], values: List[float], queries: List[List[str]]) -> List[float]:
        equationsSize = len(equations)
        uf = UnionFind(2 * equationsSize)

        # 第 1 步：预处理，将变量的值与 id 进行映射，使得并查集的底层使用数组实现，方便编码
        hashmap = dict()
        id = 0
        for i in range(equationsSize):
            equation = equations[i]
            var1 = equation[0]
            var2 = equation[1]

            if var1 not in hashmap:
                hashmap[var1] = id
                id += 1
            
            if var2 not in hashmap:
                hashmap[var2] = id
                id += 1

            uf.union(hashmap.get(var1), hashmap.get(var2), values[i])

        # 第 2 步：做查询
        queriesSize = len(queries)
        res = [None] * queriesSize
        for i in range(queriesSize):
            var1 = queries[i][0]
            var2 = queries[i][1]

            id1 = hashmap.get(var1)
            id2 = hashmap.get(var2)

            if id1 is None or id2 is None:
                res[i] = -1.0
            else:
                res[i] = uf.isConnected(id1, id2)

        return res




class UnionFind:
    """
    并查集算法
    """
    def __init__(self, n):
        self.parent = [i for i in range(n)]
        self.cnt = n
        self.weight = [1.0 for _ in range(n)]

    
    def union(self, x, y, value):
        rootx = self.find(x)
        rooty = self.find(y)
        if rootx == rooty:
            return

        self.parent[rootx] = rooty
        self.weight[rootx] = self.weight[y] * value / self.weight[x]

    def find(self, x):
        if x != self.parent[x]:
            origin = self.parent[x]
            self.parent[x] = self.find(self.parent[x])
            self.weight[x] = self.weight[x] * self.weight[origin]

        return self.parent[x]

    def isConnected(self, x, y):
        rootx = self.find(x)
        rooty = self.find(y)

        if rootx == rooty:
            return self.weight[x] / self.weight[y]
        
        return -1