added doctests to functions in lowest_common_ancestor.py

Author: NidhaNureen

data_structures/binary_tree/lowest_common_ancestor.py

@@ -25,6 +25,16 @@ def swap(a: int, b: int) -> tuple[int, int]:
 def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
     """
     creating sparse table which saves each nodes 2^i-th parent
+    >>> max_node = 6
+    >>> parent = [[0, 0, 1, 1, 2, 2, 3]] + [[0] * 7 for _ in range(19)]
+    >>> parent = create_sparse(max_node, parent)
+    >>> parent[0]
+    [0, 0, 1, 1, 2, 2, 3]
+    >>> parent[1]
+    [0, 0, 0, 0, 1, 1, 1]
+    >>> parent[2]
+    [0, 0, 0, 0, 0, 0, 0]
+
     """
     j = 1
     while (1 << j) < max_node:
@@ -38,6 +48,26 @@ def create_sparse(max_node: int, parent: list[list[int]]) -> list[list[int]]:
 def lowest_common_ancestor(
     u: int, v: int, level: list[int], parent: list[list[int]]
 ) -> int:
+    """
+    Return the lowest common ancestor between u and v
+
+    >>> level = [-1, 0, 1, 1, 2, 2, 2]
+    >>> parent = [[0, 0, 1, 1, 2, 2, 3],[0, 0, 0, 0, 1, 1, 1]] + [[0] * 7 for _ in range(17)]
+    >>> lowest_common_ancestor(4, 5, level, parent)
+    2
+    >>> lowest_common_ancestor(3, 6, level, parent)
+    3
+    >>> lowest_common_ancestor(4, 6, level, parent)
+    1
+    >>> lowest_common_ancestor(5, 6, level, parent)
+    1
+    >>> lowest_common_ancestor(2, 3, level, parent)
+    1
+    >>> lowest_common_ancestor(6, 6, level, parent)
+    6
+    >>> lowest_common_ancestor(1, 3, level, parent)
+    1
+    """
     # u must be deeper in the tree than v
     if level[u] < level[v]:
         u, v = swap(u, v)
@@ -56,6 +86,8 @@ def lowest_common_ancestor(
     return parent[0][u]
 
 
+
+
 # runs a breadth first search from root node of the tree
 def breadth_first_search(
     level: list[int],
@@ -68,6 +100,30 @@ def breadth_first_search(
     sets every nodes direct parent
     parent of root node is set to 0
     calculates depth of each node from root node
+    >>> level = [-1] * 7
+    >>> parent = [[0] * 7 for _ in range(20)]
+    >>> graph = {1: [2, 3], 2: [4, 5], 3: [6], 4: [], 5: [], 6: []}
+    >>> level, parent = breadth_first_search(level, parent, 6, graph, 1)
+    >>> level
+    [-1, 0, 1, 1, 2, 2, 2]
+    >>> parent[0]
+    [0, 0, 1, 1, 2, 2, 3]
+    >>> parent[1]
+    [0, 0, 0, 0, 0, 0, 0]
+    >>> parent[2]
+    [0, 0, 0, 0, 0, 0, 0]
+
+    # Edge case: graph with one node
+    >>> level = [-1] * 2
+    >>> parent = [[0] * 2 for _ in range(20)]
+    >>> graph = {1: []}
+    >>> level, parent = breadth_first_search(level, parent, 1, graph, 1)
+    >>> level
+    [-1, 0]
+    >>> parent[0]
+    [0, 0]
+    >>> parent[1]
+    [0, 0]
     """
     level[root] = 0
     q: Queue[int] = Queue(maxsize=max_node)