segment_tree.py

importmath


classSegmentTree:
def__init__(self, 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
ifself.N:
self.build(1, 0, self.N-1)

defleft(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
 """
returnidx*2

defright(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
 """
returnidx*2+1

defbuild(self, idx, left, right):
ifleft==right:
self.st[idx] =self.A[left]
else:
mid= (left+right) //2
self.build(self.left(idx), left, mid)
self.build(self.right(idx), mid+1, right)
self.st[idx] =max(self.st[self.left(idx)], self.st[self.right(idx)])

defupdate(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
 """
returnself.update_recursive(1, 0, self.N-1, a-1, b-1, val)

defupdate_recursive(self, idx, left, right, a, b, val):
"""
 update(1, 1, N, a, b, v) for update val v to [a,b]
 """
ifright<aorleft>b:
returnTrue
ifleft==right:
self.st[idx] =val
returnTrue
mid= (left+right) //2
self.update_recursive(self.left(idx), left, mid, a, b, val)
self.update_recursive(self.right(idx), mid+1, right, a, b, val)
self.st[idx] =max(self.st[self.left(idx)], self.st[self.right(idx)])
returnTrue

defquery(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
 """
returnself.query_recursive(1, 0, self.N-1, a-1, b-1)

defquery_recursive(self, idx, left, right, a, b):
"""
 query(1, 1, N, a, b) for query max of [a,b]
 """
ifright<aorleft>b:
return-math.inf
ifleft>=aandright<=b:
returnself.st[idx]
mid= (left+right) //2
q1=self.query_recursive(self.left(idx), left, mid, a, b)
q2=self.query_recursive(self.right(idx), mid+1, right, a, b)
returnmax(q1, q2)

defshow_data(self):
show_list= []
foriinrange(1, self.N+1):
show_list+= [self.query(i, i)]
print(show_list)


if__name__=="__main__":
A= [1, 2, -4, 7, 3, -5, 6, 11, -20, 9, 14, 15, 5, 2, -8]
N=15
segt=SegmentTree(A)
print(segt.query(4, 6))
print(segt.query(7, 11))
print(segt.query(7, 12))
segt.update(1, 3, 111)
print(segt.query(1, 15))
segt.update(7, 8, 235)
segt.show_data()