LeetCode #399: Evaluate Division

Description:

Equations are given in the format A / B = k, where A and B are variables represented as strings, and k is a real number (floating point number). Given some queries, return the answers. If the answer does not exist, return -1.0.

Example:
Given a / b = 2.0, b / c = 3.0.
queries are: a / c = ?, b / a = ?, a / e = ?, a / a = ?, x / x = ? .
return [6.0, 0.5, -1.0, 1.0, -1.0 ].

The input is: vector<pair<string, string>> equations, vector<double>& values, vector<pair<string, string>> queries , where equations.size() == values.size(), and the values are positive. This represents the equations. Return vector<double>.

According to the example above:

equations = [ ["a", "b"], ["b", "c"] ],
values = [2.0, 3.0],
queries = [ ["a", "c"], ["b", "a"], ["a", "e"], ["a", "a"], ["x", "x"] ]. 

The input is always valid. You may assume that evaluating the queries will result in no division by zero and there is no contradiction.

Solution:

• Build disjoint-set tree to restore the relations: label each element with a reference (root node)
• Time complexity: construction O(N); query O(log(N))
• Space complexity: O(N)

class Solution:
    def calcEquation(self, equations: List[List[str]], values: List[float], queries: List[List[str]]) -> List[float]:
        from collections import defaultdict

        parent = defaultdict(None)
        depth = defaultdict(int)
        wt = defaultdict(lambda: 0.0)
        for i in range(len(values)):
            if equations[i][0] in list(parent.keys()):
                p, fac = self.findparent(parent, wt, equations[i][0], 1.0 )
                parent[p] = equations[i][1]
                wt[p] = 1.0/values[i]/fac
            else:
                parent[equations[i][0]] = equations[i][1]
                wt[equations[i][0]] = 1.0/values[i]

            if equations[i][1] not in parent.keys():
                parent[equations[i][1]] = None
                wt[equations[i][1]] = 1.0

        out = []
        for q in queries:
            if q[0] not in list(parent.keys()) or q[1] not in list(parent.keys()):
                out.append(-1.0)
            else:
                px, facx = self.findparent(parent, wt, q[0], 1.0)
                py, facy = self.findparent(parent, wt, q[1], 1.0)
                if px == py:
                    out.append(facy/facx)
                else:
                    out.append(-1.0)
        return  out

    # find the root node in the union set tree
    def findparent(self, parent, wt, x, fac):
        flag = 2

        # method 1: find the root node recursively
        if flag == 1:
            if parent[x] is None:
                return x, fac
            else:
                return self.findparent(parent, wt, parent[x], fac*wt[x])

        # method 2: find the root node iteratively
        if flag == 2:
            fac = 1.0
            current = x
            while parent[current] is not None:
                fac = fac * wt[current]
                current = parent[current]
            return current, fac