木における最短経路探索 - BFS解説

from typing import List, Dict
from collections import defaultdict, deque
import sys

class Solution:
    def shortest_path_in_tree(self, n: int, x: int, y: int, edges: List[List[int]]) -> List[int]:
        """
        木における2頂点間の最短経路を BFS で求める
        """
        if not (1 <= n <= 100000):
            raise ValueError("Invalid number of nodes")
        if not (1 <= x <= n and 1 <= y <= n and x != y):
            raise ValueError("Invalid start or end node")
        if len(edges) != n - 1:
            raise ValueError("Edges must contain exactly n-1 items")

        graph: Dict[int, List[int]] = defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)

        parent = [-1] * (n + 1)
        parent[x] = 0
        q = deque([x])

        while q:
            cur = q.popleft()
            if cur == y:
                break
            for nxt in graph[cur]:
                if parent[nxt] == -1:
                    parent[nxt] = cur
                    q.append(nxt)

        path: List[int] = []
        node = y
        while node != 0:
            path.append(node)
            if node == x:
                break
            node = parent[node]
        path.reverse()
        return path

if __name__ == "__main__":
    sys.setrecursionlimit(1 << 25)
    N, X, Y = 6, 1, 6
    edges = [[1,2],[1,3],[2,4],[2,5],[3,6]]
    solver = Solution()
    print(solver.shortest_path_in_tree(N, X, Y, edges))