modified files to meet style guidelines, edited docstrings and added some doctests

Author: kanthuc

data_structures/binary_tree/lazy_segment_tree.py

@@ -5,18 +5,32 @@
 class SegmentTree:
     def __init__(self, N: int) -> None:
         self.N = N
-        self.st: List[int] = [
-            0 for i in range(0, 4 * N)
-        ]  # approximate the overall size of segment tree with array N
-        self.lazy: List[int] = [
-            0 for i in range(0, 4 * N)
-        ]  # create array to store lazy update
+        # approximate the overall size of segment tree with array N
+        self.st: List[int] = [0 for i in range(0, 4 * N)]
+        # create array to store lazy update
+        self.lazy: List[int] = [0 for i in range(0, 4 * N)]
         self.flag: List[int] = [0 for i in range(0, 4 * N)]  # flag for lazy update
 
     def left(self, idx: int) -> int:
+        """
+        >>> left(1)
+        2
+        >>> left(2)
+        4
+        >>> left(12)
+        24
+        """
         return idx * 2
 
     def right(self, idx: int) -> int:
+        """
+        >>> left(1)
+        3
+        >>> left(2)
+        5
+        >>> left(12)
+        25
+        """
         return idx * 2 + 1
 
     def build(
@@ -30,12 +44,13 @@ def build(
             self.build(self.right(idx), mid + 1, right_element, A)
             self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)])
 
-    # update with O(lg N) (Normal segment tree without lazy update will take O(Nlg N)
-    # for each update)
     def update(
         self, idx: int, left_element: int, right_element: int, a: int, b: int, val: int
     ) -> bool:
         """
+        update with O(lg N) (Normal segment tree without lazy update will take O(Nlg N)
+        for each update)
+
         update(1, 1, N, a, b, v) for update val v to [a,b]
         """
         if self.flag[idx] is True:
@@ -87,7 +102,7 @@ def query(
         q2 = self.query(self.right(idx), mid + 1, right_element, a, b)
         return max(q1, q2)
 
-    def showData(self) -> None:
+    def show_data(self) -> None:
         showList = []
         for i in range(1, N + 1):
             showList += [self.query(1, 1, self.N, i, i)]
@@ -105,4 +120,4 @@ def showData(self) -> None:
     segt.update(1, 1, N, 1, 3, 111)
     print(segt.query(1, 1, N, 1, 15))
     segt.update(1, 1, N, 7, 8, 235)
-    segt.showData()
+    segt.show_data()

data_structures/binary_tree/lowest_common_ancestor.py

@@ -2,18 +2,29 @@
 # https://en.wikipedia.org/wiki/Breadth-first_search
 
 import queue
-from typing import Tuple, List, Dict
+from typing import Dict, List, Tuple
 
 
 def swap(a: int, b: int) -> Tuple[int, int]:
+    """
+    Return a tuple (b, a) when given two integers a and b
+    >>> swap(2,3)
+    (3,2)
+    >>> swap(3,4)
+    (4,3)
+    >>> swap(67, 12)
+    (12, 67)
+    """
     a ^= b
     b ^= a
     a ^= b
     return a, b
 
 
-# creating sparse table which saves each nodes 2^i-th parent
-def creatSparse(max_node: int, parent: List[List[int]]) -> List[List[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
+    """
     j = 1
     while (1 << j) < max_node:
         for i in range(1, max_node + 1):
@@ -23,7 +34,9 @@ def creatSparse(max_node: int, parent: List[List[int]]) -> List[List[int]]:
 
 
 # returns lca of node u,v
-def LCA(u: int, v: int, level: List[int], parent: List[List[int]]) -> List[List[int]]:
+def lowest_common_ancestor(
+    u: int, v: int, level: List[int], parent: List[List[int]]
+) -> List[List[int]]:
     # u must be deeper in the tree than v
     if level[u] < level[v]:
         u, v = swap(u, v)
@@ -43,16 +56,18 @@ def LCA(u: int, v: int, level: List[int], parent: List[List[int]]) -> List[List[
 
 
 # runs a breadth first search from root node of the tree
-# sets every nodes direct parent
-# parent of root node is set to 0
-# calculates depth of each node from root node
-def bfs(
+def breadth_first_search(
     level: List[int],
     parent: List[List[int]],
     max_node: int,
     graph: Dict[int, int],
     root=1,
 ) -> Tuple[List[int], List[List[int]]]:
+    """
+    sets every nodes direct parent
+    parent of root node is set to 0
+    calculates depth of each node from root node
+    """
     level[root] = 0
     q = queue.Queue(maxsize=max_node)
     q.put(root)
@@ -87,14 +102,14 @@ def main() -> None:
         12: [],
         13: [],
     }
-    level, parent = bfs(level, parent, max_node, graph, 1)
-    parent = creatSparse(max_node, parent)
-    print("LCA of node 1 and 3 is: ", LCA(1, 3, level, parent))
-    print("LCA of node 5 and 6 is: ", LCA(5, 6, level, parent))
-    print("LCA of node 7 and 11 is: ", LCA(7, 11, level, parent))
-    print("LCA of node 6 and 7 is: ", LCA(6, 7, level, parent))
-    print("LCA of node 4 and 12 is: ", LCA(4, 12, level, parent))
-    print("LCA of node 8 and 8 is: ", LCA(8, 8, level, parent))
+    level, parent = breadth_first_search(level, parent, max_node, graph, 1)
+    parent = create_sparse(max_node, parent)
+    print("LCA of node 1 and 3 is: ", lowest_common_ancestor(1, 3, level, parent))
+    print("LCA of node 5 and 6 is: ", lowest_common_ancestor(5, 6, level, parent))
+    print("LCA of node 7 and 11 is: ", lowest_common_ancestor(7, 11, level, parent))
+    print("LCA of node 6 and 7 is: ", lowest_common_ancestor(6, 7, level, parent))
+    print("LCA of node 4 and 12 is: ", lowest_common_ancestor(4, 12, level, parent))
+    print("LCA of node 8 and 8 is: ", lowest_common_ancestor(8, 8, level, parent))
 
 
 if __name__ == "__main__":