diff --git a/data_structures/binary_tree/lowest_common_ancestor.py b/data_structures/binary_tree/lowest_common_ancestor.py index f34e6f772..938446946 100644 --- a/data_structures/binary_tree/lowest_common_ancestor.py +++ b/data_structures/binary_tree/lowest_common_ancestor.py @@ -1,6 +1,3 @@ -# https://en.wikipedia.org/wiki/Lowest_common_ancestor -# https://en.wikipedia.org/wiki/Breadth-first_search - from __future__ import annotations from queue import Queue @@ -25,29 +22,27 @@ def swap(a: int, b: int) -> tuple[int, int]: 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. - The function then fills in parent[j][i] = parent[j-1][parent[j-1][i]] for each j where 2^j < max_node. - + + 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.) - - >>> # Create an initial parent table with 2 rows and indices 0..3. - >>> parent0 = [0, 0, 1, 1] # 0 is unused; node1's parent=0, node2 and 3's parent=1. + + >>> 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] - >>> # We need at least (1 << j) < max_node holds only for j = 1 here since (1 << 1)=2 < 3 and (1 << 2)=4 !< 3. >>> sparse = create_sparse(3, parent) - >>> sparse[1][1], sparse[1][2], sparse[1][3] + >>> (sparse[1][1], sparse[1][2], sparse[1][3]) (0, 0, 0) - >>> # Explanation: - >>> # For node 1: parent[1][1] = parent[0][parent[0][1]] = parent[0][0] = 0. - >>> # For node 2: parent[1][2] = parent[0][parent[0][2]] = parent[0][1] = 0. - >>> # For node 3: parent[1][3] = parent[0][parent[0][3]] = parent[0][1] = 0. """ j = 1 while (1 << j) < max_node: @@ -62,20 +57,20 @@ def lowest_common_ancestor( ) -> int: """ Return the lowest common ancestor (LCA) of nodes u and v in a tree. - - 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. >>> # 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] + >>> # 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 - >>> # LCA of a node with itself is itself. >>> lowest_common_ancestor(2, 2, level, parent) 2 """ @@ -93,7 +88,6 @@ def lowest_common_ancestor( 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 the parent (direct ancestor) which is the LCA. return parent[0][u] @@ -106,10 +100,10 @@ def breadth_first_search( ) -> tuple[list[int], list[list[int]]]: """ Run a breadth-first search (BFS) from the root node of the tree. - - Sets every node's direct parent (in parent[0]) and calculates the depth (level) - of each node from the root. - + + 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 >>> # / \\ @@ -138,19 +132,19 @@ def breadth_first_search( 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. - + 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: - 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 + The expected lowest common ancestors are: - LCA(1, 3) --> 1 - LCA(5, 6) --> 1 @@ -158,7 +152,7 @@ def main() -> None: - 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 @@ -174,9 +168,7 @@ def main() -> None: True """ max_node = 13 - # initializing with 0; extra space is allocated. parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] - # initializing with -1 which means every node is unvisited. level = [-1 for _ in range(max_node + 10)] graph: dict[int, list[int]] = { 1: [2, 3, 4],