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

DFS directed-graph, O(nk), S(n ^ 2)

注意一定需要一个visited set记录已经访问过的节点，否则重复访问会导致死循环。

class Solution {
    public double[] calcEquation(String[][] equations, double[] values, String[][] queries) {
        Map<String, Map<String, Double>> map = new HashMap<>();
        
        for (int i = 0; i < equations.length; ++i) {
            if (!map.containsKey(equations[i][0])) {
                map.put(equations[i][0], new HashMap<>());
            }
            
            if (!map.containsKey(equations[i][1])) {
                map.put(equations[i][1], new HashMap<>());
            }
            
            map.get(equations[i][0]).put(equations[i][1], values[i]);
            map.get(equations[i][1]).put(equations[i][0], 1 / values[i]);
        }
        
        double[] res = new double[queries.length];
        Set<String> visited = new HashSet<>();
        
        for (int i = 0; i < res.length; ++i) {
            res[i] = dfsCalcEquation(queries[i][0], queries[i][1], map, visited);
        }
        
        return res;
    }
    
    public double dfsCalcEquation(
        String a, String b, Map<String, Map<String, Double>> map, Set<String> visited) {
        
        if (!map.containsKey(a) || !map.containsKey(b) || visited.contains(a)) {
            return -1;
        }
        
        if (a.equals(b)) {
            return 1;
        }
        
        visited.add(a);
        double res = -1;
        
        for (Map.Entry<String, Double> entry : map.get(a).entrySet()) {
            res = dfsCalcEquation(entry.getKey(), b, map, visited);
            if (res > 0) {
                res *= entry.getValue();
                break;
            }
        }
        
        visited.remove(a);
        return res;
    }
}

Optimized: Union-Find + Union by rank + Path compression, O(n + k), S(n)

这个思路真心牛掰，核心是`将common divisor作为root！
假设给定Graph是这样的（图左边）
新添加表达式a / e = 10.0之后，graph变成（图右边）

image.png

理解了思路之后，发现也并不难写。

class Solution {
    public double[] calcEquation(String[][] equations, double[] values, String[][] queries) {
        Map<String, Node> graph = new HashMap<>();
        // build graph
        for (int i = 0; i < equations.length; ++i) {
            String a = equations[i][0];
            String b = equations[i][1];
            double val = values[i];
            // initialize graph node if not exist
            if (!graph.containsKey(a)) {
                graph.put(a, new Node());
            }
            if (!graph.containsKey(b)) {
                graph.put(b, new Node());
            }
            // union!
            union(graph.get(a), graph.get(b), val);
        }
        
        double[] res = new double[queries.length];
        
        for (int i = 0; i < queries.length; ++i) {
            String a = queries[i][0];
            String b = queries[i][1];
            // decide if two nodes exist and are connected
            if (!graph.containsKey(a) || !graph.containsKey(b) 
                || !isConnected(graph.get(a), graph.get(b))) {
                res[i] = -1;
            } else {
                res[i] = graph.get(a).val / graph.get(b).val;
            }
        }
        
        return res;
    }
    
    private boolean isConnected(Node node1, Node node2) {
        return find(node1) == find(node2);
    }
    
    private Node find(Node node) {
        if (node.parent == null) {
            return node;
        }
        node.parent = find(node.parent);
        return node.parent;
    }
    
    private void union(Node node1, Node node2, double val) {
        if (isConnected(node1, node2)) { // important! don't union connected ones
            return;
        }
        
        Node root1 = find(node1);
        Node root2 = find(node2);
        // make sure we union smaller set to bigger set
        if (root1.children.size() <= root2.children.size()) {
            unionHelper(node1, node2, val);
        } else {
            unionHelper(node2, node1, 1 / val);
        }
    }
    // mount node1 set to node2 set, node1 is smaller set, path compression!
    private void unionHelper(Node node1, Node node2, double val) {
        Node root1 = find(node1);
        Node root2 = find(node2);
        double ratio = node2.val * val / node1.val;
        // mount root1.children to root2       
        for (Node child : root1.children) {
            child.val *= ratio;
            child.parent = root2;
            root2.children.add(child);
        }
        // don't forget to make root2 the parent of root1 too!
        root1.val = ratio;
        root1.parent = root2;
        root2.children.add(root1);
        root1.children.clear();
    }
    
    class Node {
        Node parent;
        List<Node> children;
        double val;
        
        public Node() {
            val = 1d;   // because x / x = 1
            children = new ArrayList<>();
        }
    }
}