LeetCode 399. Evaluate Division

December 06, 2020

We can represent our input system of variables with a weighted directed graph.

If you walk from one node (variable) to another and multiply edge weights along the way, each step cancels out the intermediate node and gives you the answer to <start of walk> / <end of walk>.

e.g. multiplying $a \rarr b \rarr c$ gives us $\frac{a}{b} * \frac{b}{c} = \frac{a}{c}$

The runtime of this algorithm is $O(queries * equations^2)$

from collections import defaultdict
from typing import List
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertEqual(self.s.calcEquation([["a","b"],["b","c"]], [2.0,3.0], [["a","c"],["b","a"],["a","e"],["a","a"],["x","x"]]), [6.00000,0.50000,-1.00000,1.00000,-1.00000])
        self.assertEqual(self.s.calcEquation([["a","b"],["b","c"],["bc","cd"]], [1.5,2.5,5.0], [["a","c"],["c","b"],["bc","cd"],["cd","bc"]]), [3.75000,0.40000,5.00000,0.20000])
        self.assertEqual(self.s.calcEquation([["a","b"]], [0.5], [["a","b"],["b","a"],["a","c"],["x","y"]]), [0.50000,2.00000,-1.00000,-1.00000])

class Solution:

    def _dfs_divide(self, start, end, graph, seen, product=1.0):
        if start not in graph or start in seen:
            return None
        if start == end:
            return product
        seen.add(start)
        for n, v in graph[start].items():
            possibility = self._dfs_divide(n, end, graph, seen, product*v)
            if possibility != None:
                return possibility
        return None

    def calcEquation(self, equations: List[List[str]], values: List[float], queries: List[List[str]]) -> List[float]:
        answers = []
        edge_list = defaultdict(lambda: defaultdict(int))
        for eq, v in zip(equations, values):
            edge_list[eq[0]][eq[1]] = v
            edge_list[eq[1]][eq[0]] = 1/v
        for q in queries:
            # Do a graph search / walk from q[0] -> q[1]
            # Use a DFS with a short-circuit
            answers.append(self._dfs_divide(q[0], q[1], edge_list, set()))
        return [-1.0 if x == None else x for x in answers]

if __name__ == "__main__":
    unittest.main()