From b74d7f53925cee24bddd436d6a7d6766d573a17e Mon Sep 17 00:00:00 2001 From: Siddhant Jain Date: Mon, 13 Jan 2025 16:56:22 -0500 Subject: [PATCH] Modified doctest --- .../binary_tree/lowest_common_ancestor.py | 210 +++++++++--------- 1 file changed, 100 insertions(+), 110 deletions(-) diff --git a/data_structures/binary_tree/lowest_common_ancestor.py b/data_structures/binary_tree/lowest_common_ancestor.py index 830f3f85c..f34e6f772 100644 --- a/data_structures/binary_tree/lowest_common_ancestor.py +++ b/data_structures/binary_tree/lowest_common_ancestor.py @@ -2,16 +2,16 @@ # https://en.wikipedia.org/wiki/Breadth-first_search 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) + Return a tuple (b, a) when given two integers a and b. + + >>> swap(2, 3) (3, 2) - >>> swap(3,4) + >>> swap(3, 4) (4, 3) >>> swap(67, 12) (12, 67) @@ -24,22 +24,30 @@ 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 which saves each node's 2^i-th parent. + Create a sparse table that saves each node's 2^i-th parent. - >>> max_node = 5 - >>> parent = [ - ... [0, 0, 1, 1, 2, 2], # 2^0-th parents - ... [0, 0, 0, 0, 1, 1] # 2^1-th parents - ... ] - >>> create_sparse(max_node, parent) - [[0, 0, 1, 1, 2, 2], [0, 0, 0, 0, 1, 1]] - >>> max_node = 3 - >>> parent = [ - ... [0, 0, 1, 1], # 2^0-th parents - ... [0, 0, 0, 0] # 2^1-th parents - ... ] - >>> create_sparse(max_node, parent) - [[0, 0, 1, 1], [0, 0, 0, 0]] + 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. + + 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. + >>> 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] + (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: @@ -49,69 +57,46 @@ def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]: return parent -# returns lca of node u,v def lowest_common_ancestor( u: int, v: int, level: list[int], parent: list[list[int]] ) -> int: """ - Return the lowest common ancestor of nodes u and v. + Return the lowest common ancestor (LCA) of nodes u and v in a tree. - >>> max_node = 13 - >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] - >>> level = [-1 for _ in range(max_node + 10)] - >>> graph = { - ... 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) - >>> lowest_common_ancestor(1, 3, level, parent) + 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] + >>> 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(5, 6, level, parent) - 1 - >>> lowest_common_ancestor(7, 11, level, parent) - 1 - >>> lowest_common_ancestor(6, 7, level, parent) - 3 - >>> lowest_common_ancestor(4, 12, level, parent) - 4 - >>> lowest_common_ancestor(8, 8, level, parent) - 8 - >>> lowest_common_ancestor(9, 10, level, parent) - 5 - >>> lowest_common_ancestor(12, 13, level, parent) - 8 + >>> # LCA of a node with itself is itself. + >>> lowest_common_ancestor(2, 2, level, parent) + 2 """ - # u must be deeper in the tree than v + # Ensure u is at least as deep as v. if level[u] < level[v]: u, v = swap(u, v) - # making depth of u same as depth of 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] - # at the same depth if u==v that mean lca is found + # If they are the same, we've found the LCA. if u == v: return u - # moving both nodes upwards till lca in found + # 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] - # returning longest common ancestor of u,v + # Return the parent (direct ancestor) which is the LCA. return parent[0][u] -# runs a breadth first search from root node of the tree def breadth_first_search( level: list[int], parent: list[list[int]], @@ -120,54 +105,23 @@ def breadth_first_search( root: int = 1, ) -> tuple[list[int], list[list[int]]]: """ - Perform a breadth-first search from the root node of the tree. - Sets every node's direct parent and calculates the depth of each node from the root. + Run a breadth-first search (BFS) from the root node of the tree. - >>> max_node = 5 - >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] - >>> level = [-1 for _ in range(max_node + 10)] - >>> graph = { - ... 1: [2, 3], - ... 2: [4], - ... 3: [5], - ... 4: [], - ... 5: [] - ... } - >>> level, parent = breadth_first_search(level, parent, max_node, graph, 1) - >>> level[:6] - [ -1, 0, 1, 1, 2, 2] - >>> parent[0][1] == 0 - True - >>> parent[0][2] == 1 - True - >>> parent[0][3] == 1 - True - >>> parent[0][4] == 2 - True - >>> parent[0][5] == 3 - True + Sets every node's direct parent (in parent[0]) and calculates the depth (level) + of each node from the root. - >>> # Test with disconnected graph - >>> max_node = 4 - >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] - >>> level = [-1 for _ in range(max_node + 10)] - >>> graph = { - ... 1: [2], - ... 2: [], - ... 3: [4], - ... 4: [] - ... } - >>> level, parent = breadth_first_search(level, parent, max_node, graph, 1) - >>> level[:5] - [ -1, 0, 1, -1, -1] - >>> parent[0][1] == 0 - True - >>> parent[0][2] == 1 - True - >>> parent[0][3] == 0 - True - >>> parent[0][4] == 3 - True + >>> # 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) @@ -183,10 +137,46 @@ 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. + + The sample tree used is: + + 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 + - 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 - # initializing with 0 + # 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 + # 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],