Update closest-node-to-path-in-tree.py

Author: kamyu104

Python/closest-node-to-path-in-tree.py

@@ -205,7 +205,7 @@ def divide(curr, parent):
             stk.pop()()
         self.D, self.P = D, P
 
-    def lca(self, a, b):  # Time: O(H)
+    def lca(self, a, b):  # Time: O(logh)
         while self.D[a] > self.D[b]:
             a = self.P[a]
         while self.D[a] < self.D[b]:
@@ -229,3 +229,38 @@ def closestNode(self, n, edges, query):
             adj[u].append(v), adj[v].append(u)
         tree_infos = TreeInfos3(adj)
         return [max((tree_infos.lca(x, y) for x, y in ((start, end), (start, node), (end, node))), key=lambda x: tree_infos.D[x]) for start, end, node in query]
+
+
+# Time:  O(n^2 + q * n)
+# Space: O(n^2)
+# bfs
+class Solution4(object):
+    def closestNode(self, n, edges, query):
+        """
+        :type n: int
+        :type edges: List[List[int]]
+        :type query: List[List[int]]
+        :rtype: List[int]
+        """
+        def bfs(adj, root):
+            dist = [len(adj)]*len(adj)
+            q = [root]
+            dist[root] = 0
+            d = 0
+            while q:
+                new_q = []
+                for u in q:
+                    for v in adj[u]:
+                        if d+1 >= dist[v]:
+                            continue
+                        dist[v] = d+1
+                        new_q.append(v)
+                q = new_q
+                d += 1
+            return dist
+
+        adj = [[] for _ in xrange(n)]
+        for u, v in edges:
+            adj[u].append(v), adj[v].append(u)
+        dist = [bfs(adj, i) for i in xrange(n)]
+        return [max((i for i in xrange(n) if dist[start][node]+dist[node][end]-2*dist[node][i] == dist[start][i]+dist[i][end]), key=lambda x: dist[node][x]) for start, end, node in query]