Modified doctest

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 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.
+    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, nodes 2 and 3's parent=1.
-    >>> 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]
+    >>> (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,
+    The lists ``level`` and ``parent`` must be precomputed. ``level[i]`` is the depth
-    and `parent` is a sparse table where parent[0][i] is the direct parent of node i.
+    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)
+    This sets each node's direct parent (stored in parent[0]) and calculates the
-    of each node from the root.
+    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,
+    Run a BFS to set node depths and parents in a sample tree, then create the
-    then create the sparse table and compute several lowest common ancestors.
+    sparse table and compute several lowest common ancestors.
-
+    
     The sample tree used is:
     
-            1
+                1
-         /  |  \ 
+             /  |  \ 
-        2   3   4
+            2   3   4
-       /   / \\  \\
+           /   / \\   \\
-      5   6   7  8
+          5   6   7   8
-     / \\  |     / \\
+         / \\   |    / \\
-    9  10  11   12 13
+        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],