From 8b1de765fff93e11e3ead7ac523e35dba5d5eec4 Mon Sep 17 00:00:00 2001 From: Aryan Gandhi Date: Sat, 7 Oct 2023 00:15:55 -0500 Subject: [PATCH 1/3] update docstrings --- data_structures/binary_tree/segment_tree.py | 36 +++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/data_structures/binary_tree/segment_tree.py b/data_structures/binary_tree/segment_tree.py index 5f822407d8cb..a3bb7e1ee1f9 100644 --- a/data_structures/binary_tree/segment_tree.py +++ b/data_structures/binary_tree/segment_tree.py @@ -11,9 +11,27 @@ def __init__(self, a): self.build(1, 0, self.N - 1) def left(self, idx): + """ + Returns the left child index for a given index in a binary tree. + + >>> s = SegmentTree([1, 2, 3]) + >>> s.left(1) + 2 + >>> s.left(2) + 4 + """ return idx * 2 def right(self, idx): + """ + Returns the right child index for a given index in a binary tree. + + >>> s = SegmentTree([1, 2, 3]) + >>> s.right(1) + 3 + >>> s.right(2) + 5 + """ return idx * 2 + 1 def build(self, idx, l, r): # noqa: E741 @@ -26,6 +44,15 @@ def build(self, idx, l, r): # noqa: E741 self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)]) def update(self, a, b, val): + """ + Update the values in the segment tree in the range [a,b] with the given value. + + >>> s = SegmentTree([1, 2, 3, 4, 5]) + >>> s.update(2, 4, 10) + True + >>> s.query(1, 5) + 10 + """ return self.update_recursive(1, 0, self.N - 1, a - 1, b - 1, val) def update_recursive(self, idx, l, r, a, b, val): # noqa: E741 @@ -44,6 +71,15 @@ def update_recursive(self, idx, l, r, a, b, val): # noqa: E741 return True def query(self, a, b): + """ + Query the maximum value in the range [a,b]. + + >>> s = SegmentTree([1, 2, 3, 4, 5]) + >>> s.query(1, 3) + 3 + >>> s.query(1, 5) + 5 + """ return self.query_recursive(1, 0, self.N - 1, a - 1, b - 1) def query_recursive(self, idx, l, r, a, b): # noqa: E741 From feaca465cb78df74e5b7ae005c3b3fb265a1865d Mon Sep 17 00:00:00 2001 From: Aryan Gandhi Date: Sat, 7 Oct 2023 00:30:44 -0500 Subject: [PATCH 2/3] update docstrings --- data_structures/binary_tree/segment_tree.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/data_structures/binary_tree/segment_tree.py b/data_structures/binary_tree/segment_tree.py index a3bb7e1ee1f9..62a53bd989a5 100644 --- a/data_structures/binary_tree/segment_tree.py +++ b/data_structures/binary_tree/segment_tree.py @@ -76,9 +76,9 @@ def query(self, a, b): >>> s = SegmentTree([1, 2, 3, 4, 5]) >>> s.query(1, 3) - 3 + 2 >>> s.query(1, 5) - 5 + 7 """ return self.query_recursive(1, 0, self.N - 1, a - 1, b - 1) From a725bedead2a39876f90fad4b0d52fcc4de77b2f Mon Sep 17 00:00:00 2001 From: Aryan Gandhi Date: Sat, 7 Oct 2023 00:47:51 -0500 Subject: [PATCH 3/3] update docstrings --- data_structures/binary_tree/segment_tree.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/data_structures/binary_tree/segment_tree.py b/data_structures/binary_tree/segment_tree.py index 62a53bd989a5..3b0b32946f6e 100644 --- a/data_structures/binary_tree/segment_tree.py +++ b/data_structures/binary_tree/segment_tree.py @@ -3,7 +3,8 @@ class SegmentTree: def __init__(self, a): - self.N = len(a) + self.A = a + self.N = len(self.A) self.st = [0] * ( 4 * self.N ) # approximate the overall size of segment tree with array N @@ -36,7 +37,7 @@ def right(self, idx): def build(self, idx, l, r): # noqa: E741 if l == r: - self.st[idx] = A[l] + self.st[idx] = self.A[l] else: mid = (l + r) // 2 self.build(self.left(idx), l, mid) @@ -76,9 +77,9 @@ def query(self, a, b): >>> s = SegmentTree([1, 2, 3, 4, 5]) >>> s.query(1, 3) - 2 + 3 >>> s.query(1, 5) - 7 + 5 """ return self.query_recursive(1, 0, self.N - 1, a - 1, b - 1)