from __future__ import annotations from queue import Queue def swap(a: int, b: int) -> tuple[int, int]: """ Return a tuple (b, a) when given two integers a and b. >>> swap(2, 3) (3, 2) >>> swap(3, 4) (4, 3) >>> swap(67, 12) (12, 67) """ a ^= b b ^= a a ^= b return a, b def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]: """ Create a sparse table that saves each node's 2^i-th parent. The given ``parent`` table should have the direct parent of each node in row 0. This function fills in: parent[j][i] = parent[j - 1][parent[j - 1][i]] for each j where 2^j is less than max_node. For example, consider a small tree where: - Node 1 is the root (its parent is 0), - Nodes 2 and 3 have parent 1. We set up the parent table for only two levels (row 0 and row 1) for max_node = 3. (Note that in practice the table has many rows.) >>> parent0 = [0, 0, 1, 1] # 0 is unused; node1's parent=0, nodes 2 and 3's parent=1. >>> parent1 = [0, 0, 0, 0] >>> parent = [parent0, parent1] >>> sparse = create_sparse(3, parent) >>> (sparse[1][1], sparse[1][2], sparse[1][3]) (0, 0, 0) """ j = 1 while (1 << j) < max_node: for i in range(1, max_node + 1): parent[j][i] = parent[j - 1][parent[j - 1][i]] j += 1 return parent def lowest_common_ancestor( u: int, v: int, level: list[int], parent: list[list[int]] ) -> int: """ Return the lowest common ancestor (LCA) of nodes u and v in a tree. <<<<<<< HEAD The lists ``level`` and ``parent`` must be precomputed. ``level[i]`` is the depth of node i, and ``parent`` is a sparse table where parent[0][i] is the direct parent of node i. ======= The lists `level` and `parent` must be precomputed. `level[i]` is the depth of node i, and `parent` is a sparse table where parent[0][i] is the direct parent of node i. >>>>>>> 097e9c6149e80f095be1b3dbef1c04ff94a7325a >>> # Consider a simple tree: >>> # 1 >>> # / \\ >>> # 2 3 >>> # With levels: level[1]=0, level[2]=1, level[3]=1 and parent[0]=[0, 0, 1, 1] >>> level = [-1, 0, 1, 1] # index 0 is dummy >>> parent = [[0, 0, 1, 1]] + [[0, 0, 0, 0] for _ in range(19)] >>> lowest_common_ancestor(2, 3, level, parent) 1 >>> lowest_common_ancestor(2, 2, level, parent) 2 """ # Ensure u is at least as deep as v. if level[u] < level[v]: u, v = swap(u, v) # Bring u up to the same level as v. for i in range(18, -1, -1): if level[u] - (1 << i) >= level[v]: u = parent[i][u] # If they are the same, we've found the LCA. if u == v: return u # Move u and v up together until the LCA is found. for i in range(18, -1, -1): if parent[i][u] not in [0, parent[i][v]]: u, v = parent[i][u], parent[i][v] return parent[0][u] def breadth_first_search( level: list[int], parent: list[list[int]], max_node: int, graph: dict[int, list[int]], root: int = 1, ) -> tuple[list[int], list[list[int]]]: """ Run a breadth-first search (BFS) from the root node of the tree. This sets each node's direct parent (stored in parent[0]) and calculates the depth (level) of each node from the root. >>> # Consider a simple tree: >>> # 1 >>> # / \\ >>> # 2 3 >>> graph = {1: [2, 3], 2: [], 3: []} >>> level = [-1] * 4 # index 0 is unused; nodes 1 to 3. >>> parent = [[0] * 4 for _ in range(20)] >>> new_level, new_parent = breadth_first_search(level, parent, 3, graph, root=1) >>> new_level[1:4] [0, 1, 1] >>> new_parent[0][1:4] [0, 1, 1] """ level[root] = 0 q: Queue[int] = Queue(maxsize=max_node) q.put(root) while q.qsize() != 0: u = q.get() for v in graph[u]: if level[v] == -1: level[v] = level[u] + 1 q.put(v) parent[0][v] = u return level, parent def main() -> None: """ Run a BFS to set node depths and parents in a sample tree, then create the sparse table and compute several lowest common ancestors. The sample tree used is: <<<<<<< HEAD 1 / | \ 2 3 4 / / \\ \\ 5 6 7 8 / \\ | / \\ 9 10 11 12 13 ======= 1 / | \ 2 3 4 / / \\ \\ 5 6 7 8 / \\ | / \\ 9 10 11 12 13 >>>>>>> 097e9c6149e80f095be1b3dbef1c04ff94a7325a The expected lowest common ancestors are: - LCA(1, 3) --> 1 - LCA(5, 6) --> 1 - LCA(7, 11) --> 3 - LCA(6, 7) --> 3 - LCA(4, 12) --> 4 - LCA(8, 8) --> 8 To test main() without it printing to the console, we capture the output. >>> import sys >>> from io import StringIO >>> backup = sys.stdout >>> sys.stdout = StringIO() >>> main() >>> output = sys.stdout.getvalue() >>> sys.stdout = backup >>> 'LCA of node 1 and 3 is: 1' in output True >>> 'LCA of node 7 and 11 is: 3' in output True """ max_node = 13 parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] level = [-1 for _ in range(max_node + 10)] graph: dict[int, list[int]] = { 1: [2, 3, 4], 2: [5], 3: [6, 7], 4: [8], 5: [9, 10], 6: [11], 7: [], 8: [12, 13], 9: [], 10: [], 11: [], 12: [], 13: [], } level, parent = breadth_first_search(level, parent, max_node, graph, 1) parent = create_sparse(max_node, parent) print("LCA of node 1 and 3 is: ", lowest_common_ancestor(1, 3, level, parent)) print("LCA of node 5 and 6 is: ", lowest_common_ancestor(5, 6, level, parent)) print("LCA of node 7 and 11 is: ", lowest_common_ancestor(7, 11, level, parent)) print("LCA of node 6 and 7 is: ", lowest_common_ancestor(6, 7, level, parent)) print("LCA of node 4 and 12 is: ", lowest_common_ancestor(4, 12, level, parent)) print("LCA of node 8 and 8 is: ", lowest_common_ancestor(8, 8, level, parent)) if __name__ == "__main__": main()