Modified doctest

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
+    Return a tuple (b, a) when given two integers a and b.
-    >>> swap(2,3)
+
+    >>> 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
+    The given `parent` table should have the direct parent of each node in row 0.
-    >>> parent = [
+    The function then fills in parent[j][i] = parent[j-1][parent[j-1][i]] for each j where 2^j < max_node.
-    ...     [0, 0, 1, 1, 2, 2],  # 2^0-th parents
+
-    ...     [0, 0, 0, 0, 1, 1]   # 2^1-th parents
+    For example, consider a small tree where:
-    ... ]
+      - Node 1 is the root (its parent is 0),
-    >>> create_sparse(max_node, parent)
+      - Nodes 2 and 3 have parent 1.
-    [[0, 0, 1, 1, 2, 2], [0, 0, 0, 0, 1, 1]]
+    
-    >>> max_node = 3
+    We set up the parent table for only two levels (row 0 and row 1)
-    >>> parent = [
+    for max_node = 3. (Note that in practice the table has many rows.)
-    ...     [0, 0, 1, 1],  # 2^0-th parents
+
-    ...     [0, 0, 0, 0]   # 2^1-th parents
+    >>> # 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.
-    >>> create_sparse(max_node, parent)
+    >>> parent1 = [0, 0, 0, 0]
-    [[0, 0, 1, 1], [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
+    The lists `level` and `parent` must be precomputed. `level[i]` is the depth of node i,
-    >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
+    and `parent` is a sparse table where parent[0][i] is the direct parent of node i.
-    >>> level = [-1 for _ in range(max_node + 10)]
+    
-    >>> graph = {
+    >>> # Consider a simple tree:
-    ...     1: [2, 3, 4],
+    >>> #       1
-    ...     2: [5],
+    >>> #      / \\
-    ...     3: [6, 7],
+    >>> #     2   3
-    ...     4: [8],
+    >>> # With levels: level[1]=0, level[2]=1, level[3]=1 and parent[0]=[0,0,1,1]
-    ...     5: [9, 10],
+    >>> level = [-1, 0, 1, 1]  # index 0 is dummy
-    ...     6: [11],
+    >>> parent = [[0, 0, 1, 1]] + [[0, 0, 0, 0] for _ in range(19)]
-    ...     7: [],
+    >>> lowest_common_ancestor(2, 3, level, parent)
-    ...     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)
     1
-    >>> lowest_common_ancestor(5, 6, level, parent)
+    >>> # LCA of a node with itself is itself.
-    1
+    >>> lowest_common_ancestor(2, 2, level, parent)
-    >>> lowest_common_ancestor(7, 11, level, parent)
+    2
-    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
     """
-    # 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.
+    Run a breadth-first search (BFS) from the root node of the tree.
-    Sets every node's direct parent and calculates the depth of each node from the root.
 
-    >>> max_node = 5
+    Sets every node's direct parent (in parent[0]) and calculates the depth (level)
-    >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
+    of each node from the root.
-    >>> 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
 
-    >>> # Test with disconnected graph
+    >>> # Consider a simple tree:
-    >>> max_node = 4
+    >>> #    1
-    >>> parent = [[0 for _ in range(max_node + 10)] for _ in range(20)]
+    >>> #   / \\
-    >>> level = [-1 for _ in range(max_node + 10)]
+    >>> #  2   3
-    >>> graph = {
+    >>> graph = {1: [2, 3], 2: [], 3: []}
-    ...     1: [2],
+    >>> level = [-1] * 4   # index 0 is unused; nodes 1 to 3.
-    ...     2: [],
+    >>> parent = [[0] * 4 for _ in range(20)]
-    ...     3: [4],
+    >>> new_level, new_parent = breadth_first_search(level, parent, 3, graph, root=1)
-    ...     4: []
+    >>> new_level[1:4]
-    ... }
+    [0, 1, 1]
-    >>> level, parent = breadth_first_search(level, parent, max_node, graph, 1)
+    >>> new_parent[0][1:4]
-    >>> level[:5]
+    [0, 1, 1]
-    [ -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
     """
     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],