diff --git a/data_structures/binary_tree/lazy_segment_tree.py b/data_structures/binary_tree/lazy_segment_tree.py index 461996b87..66b995fa1 100644 --- a/data_structures/binary_tree/lazy_segment_tree.py +++ b/data_structures/binary_tree/lazy_segment_tree.py @@ -1,84 +1,119 @@ import math +from typing import List class SegmentTree: - def __init__(self, N): + def __init__(self, N: int) -> None: self.N = N - self.st = [ - 0 for i in range(0, 4 * N) - ] # approximate the overall size of segment tree with array N - self.lazy = [0 for i in range(0, 4 * N)] # create array to store lazy update - self.flag = [0 for i in range(0, 4 * N)] # flag for 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): + def left(self, idx: int) -> int: + """ + >>> segment_tree = SegmentTree(15) + >>> segment_tree.left(1) + 2 + >>> segment_tree.left(2) + 4 + >>> segment_tree.left(12) + 24 + """ return idx * 2 - def right(self, idx): + def right(self, idx: int) -> int: + """ + >>> segment_tree = SegmentTree(15) + >>> segment_tree.right(1) + 3 + >>> segment_tree.right(2) + 5 + >>> segment_tree.right(12) + 25 + """ return idx * 2 + 1 - def build(self, idx, l, r, A): # noqa: E741 - if l == r: # noqa: E741 - self.st[idx] = A[l - 1] + def build( + self, idx: int, left_element: int, right_element: int, A: List[int] + ) -> None: + if left_element == right_element: + self.st[idx] = A[left_element - 1] else: - mid = (l + r) // 2 - self.build(self.left(idx), l, mid, A) - self.build(self.right(idx), mid + 1, r, A) + mid = (left_element + right_element) // 2 + self.build(self.left(idx), left_element, mid, A) + 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, l, r, a, b, val): # noqa: E741 + 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: self.st[idx] = self.lazy[idx] self.flag[idx] = False - if l != r: # noqa: E741 + if left_element != right_element: self.lazy[self.left(idx)] = self.lazy[idx] self.lazy[self.right(idx)] = self.lazy[idx] self.flag[self.left(idx)] = True self.flag[self.right(idx)] = True - if r < a or l > b: + if right_element < a or left_element > b: return True - if l >= a and r <= b: # noqa: E741 + if left_element >= a and right_element <= b: self.st[idx] = val - if l != r: # noqa: E741 + if left_element != right_element: self.lazy[self.left(idx)] = val self.lazy[self.right(idx)] = val self.flag[self.left(idx)] = True self.flag[self.right(idx)] = True return True - mid = (l + r) // 2 - self.update(self.left(idx), l, mid, a, b, val) - self.update(self.right(idx), mid + 1, r, a, b, val) + mid = (left_element + right_element) // 2 + self.update(self.left(idx), left_element, mid, a, b, val) + self.update(self.right(idx), mid + 1, right_element, a, b, val) self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)]) return True # query with O(lg N) - def query(self, idx, l, r, a, b): # noqa: E741 + def query( + self, idx: int, left_element: int, right_element: int, a: int, b: int + ) -> int: """ query(1, 1, N, a, b) for query max of [a,b] + >>> A = [1, 2, -4, 7, 3, -5, 6, 11, -20, 9, 14, 15, 5, 2, -8] + >>> segment_tree = SegmentTree(15) + >>> segment_tree.build(1, 1, 15, A) + >>> segment_tree.query(1, 1, 15, 4, 6) + 7 + >>> segment_tree.query(1, 1, 15, 7, 11) + 14 + >>> segment_tree.query(1, 1, 15, 7, 12) + 15 """ if self.flag[idx] is True: self.st[idx] = self.lazy[idx] self.flag[idx] = False - if l != r: # noqa: E741 + if left_element != right_element: self.lazy[self.left(idx)] = self.lazy[idx] self.lazy[self.right(idx)] = self.lazy[idx] self.flag[self.left(idx)] = True self.flag[self.right(idx)] = True - if r < a or l > b: + if right_element < a or left_element > b: return -math.inf - if l >= a and r <= b: # noqa: E741 + if left_element >= a and right_element <= b: return self.st[idx] - mid = (l + r) // 2 - q1 = self.query(self.left(idx), l, mid, a, b) - q2 = self.query(self.right(idx), mid + 1, r, a, b) + mid = (left_element + right_element) // 2 + q1 = self.query(self.left(idx), left_element, mid, a, b) + q2 = self.query(self.right(idx), mid + 1, right_element, a, b) return max(q1, q2) - def showData(self): + def show_data(self) -> None: showList = [] for i in range(1, N + 1): showList += [self.query(1, 1, self.N, i, i)] @@ -96,4 +131,4 @@ if __name__ == "__main__": 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() diff --git a/data_structures/binary_tree/lowest_common_ancestor.py b/data_structures/binary_tree/lowest_common_ancestor.py index f560eaa5e..c25536cda 100644 --- a/data_structures/binary_tree/lowest_common_ancestor.py +++ b/data_structures/binary_tree/lowest_common_ancestor.py @@ -2,17 +2,29 @@ # https://en.wikipedia.org/wiki/Breadth-first_search import queue +from typing import Dict, List, Tuple -def swap(a, b): +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, parent): +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): @@ -22,7 +34,9 @@ def creatSparse(max_node, parent): # returns lca of node u,v -def LCA(u, v, level, parent): +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) @@ -42,10 +56,18 @@ def LCA(u, v, level, parent): # 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(level, parent, max_node, graph, root=1): +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) @@ -59,7 +81,7 @@ def bfs(level, parent, max_node, graph, root=1): return level, parent -def main(): +def main() -> None: max_node = 13 # initializing with 0 parent = [[0 for _ in range(max_node + 10)] for _ in range(20)] @@ -80,14 +102,14 @@ def main(): 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__":