Refactor: Quy nguyen bfs for his girlfriend

Blue/Session 07 - DFS/Hackerearth_bishu-and-his-girlfriend.py

@@ -1,54 +1,42 @@
-#  Problem from Hackerearth
-#  https://www.hackerearth.com/fr/practice/algorithms/graphs/depth-first-search/practice-problems/algorithm/bishu-and-his-girlfriend/description/
+#  Quy Nguyen's late-night adventure.
 
+# https://www.hackerearth.com/fr/practice/algorithms/graphs/depth-first-search/practice-problems/algorithm/bishu-and-his-girlfriend/description/
+# Quy Nguyen lives in node 1, desperate and lonely
+# He must find the nearest "girlfriend" in the Geylang tree
+# (We all know who Quy Nguyen is really looking for)
+# If tie on distance, pick the one with smallest node number
 
-def find_girl(N, graph, girls):
-    s = []
-    visited = [False for i in range(N + 1)]
+from collections import deque
 
-    roads = [0 for i in range(N + 1)]
-
-    visited[1] = True
-    s.append(1)
-
-    selected_girl = N
-    selected_road = N
-
-    while len(s) > 0:
-        u = s[-1]
-        s.pop()
-
-        for v in graph[u]:
-            if not visited[v]:
-                roads[v] = roads[u] + 1
-                if roads[v] <= selected_road:
-                    if girls[v]:
-                        if roads[v] < selected_road or (roads[v] == selected_road and selected_girl > v):
-                            selected_girl = v
-                            selected_road = roads[v]
-                            visited[v] = True
-                    else:
-                        visited[v] = True
-                        s.append(v)
-
-    return selected_girl
-
-
-def solution():
+def solve():
     N = int(input())
-    roads = []
-    graph = [[] for i in range(N + 1)]
-    for i in range(N-1):
+    g = [[] for _ in range(N + 1)]
+    for _ in range(N - 1):
         u, v = map(int, input().split())
-        graph[u].append(v)
-        graph[v].append(u)
+        g[u].append(v)
+        g[v].append(u)
 
     Q = int(input())
-    girls = [0 for i in range(N + 1)]
-    for i in range(Q):
-        girls[int(input())] = 1
-
-    print(find_girl(N, graph, girls))
-
 
-solution()
+    # Mark the "ladies of the night" locations
+    hoes = set(int(input()) for _ in range(Q))
+
+    # BFS from Quy Nguyen's house (node 1) — the horny search algorithm
+    dist = [-1] * (N + 1)
+    dist[1] = 0
+    q = deque([1])
+    best, best_d = N, N  # worst case: the farthest hooker
+    while q:
+        u = q.popleft()
+        if u in hoes and (dist[u] < best_d or (dist[u] == best_d and u < best)):
+            best, best_d = u, dist[u]  # found a closer thot
+        for v in g[u]:
+            if dist[v] == -1:
+                dist[v] = dist[u] + 1
+                if dist[v] <= best_d:  # pruning: no point going deeper than current best
+                    q.append(v)
+
+    # Do we even have Jenny in the testcases? Wondering
+    print(best)
+
+solve()