mathdatech/Python
1
0
Fork 0
mirror of https://github.com/TheAlgorithms/Python.git synced 2025-02-25 10:28:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

update segmenttree docstrings Fixes #9943 (#9975)

Browse Source 
* update docstrings

* update docstrings

* update docstrings
This commit is contained in:
aryandgandhi 2023-10-16 20:25:07 -05:00 committed by GitHub
parent 5f629b6049
commit b5786c87d8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 39 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
data_structures/binary_tree/segment_tree.py
View File

@ -3,7 +3,8 @@ import math
class SegmentTree: class SegmentTree:
def __init__(self, a): def __init__(self, a):
self.N = len(a) self.A = a
self.N = len(self.A)
self.st = [0] * ( self.st = [0] * (
4 * self.N 4 * self.N
) # approximate the overall size of segment tree with array N ) # approximate the overall size of segment tree with array N
@ -11,14 +12,32 @@ class SegmentTree:
self.build(1, 0, self.N - 1) self.build(1, 0, self.N - 1)
def left(self, idx): 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 return idx * 2
def right(self, idx): 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 return idx * 2 + 1
def build(self, idx, l, r): # noqa: E741 def build(self, idx, l, r): # noqa: E741
if l == r: if l == r:
self.st[idx] = A[l] self.st[idx] = self.A[l]
else: else:
mid = (l + r) // 2 mid = (l + r) // 2
self.build(self.left(idx), l, mid) self.build(self.left(idx), l, mid)
@ -26,6 +45,15 @@ class SegmentTree:
self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)]) self.st[idx] = max(self.st[self.left(idx)], self.st[self.right(idx)])
def update(self, a, b, val): 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) 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 def update_recursive(self, idx, l, r, a, b, val): # noqa: E741
@ -44,6 +72,15 @@ class SegmentTree:
return True return True
def query(self, a, b): 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) return self.query_recursive(1, 0, self.N - 1, a - 1, b - 1)
def query_recursive(self, idx, l, r, a, b): # noqa: E741 def query_recursive(self, idx, l, r, a, b): # noqa: E741