Python/data_structures/binary_tree/segment_tree.py at master · FrogBattle/Python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
from __future__ importprint_function
importmath

classSegmentTree:

def__init__(self, A):
self.N=len(A)
self.st= [0] * (4*self.N) # approximate the overall size of segment tree with array N
self.build(1, 0, self.N-1)

defleft(self, idx):
returnidx*2

defright(self, idx):
returnidx*2+1

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

defupdate(self, a, b, val):
returnself.update_recursive(1, 0, self.N-1, a-1, b-1, val)

defupdate_recursive(self, idx, l, r, a, b, val): # update(1, 1, N, a, b, v) for update val v to [a,b]
ifr<aorl>b:
returnTrue
ifl==r :
self.st[idx] =val
returnTrue
mid= (l+r)//2
self.update_recursive(self.left(idx), l, mid, a, b, val)
self.update_recursive(self.right(idx), mid+1, r, a, b, val)
self.st[idx] =max(self.st[self.left(idx)] , self.st[self.right(idx)])
returnTrue

defquery(self, a, b):
returnself.query_recursive(1, 0, self.N-1, a-1, b-1)

defquery_recursive(self, idx, l, r, a, b): #query(1, 1, N, a, b) for query max of [a,b]
ifr<aorl>b:
return-math.inf
ifl>=aandr<=b:
returnself.st[idx]
mid= (l+r)//2
q1=self.query_recursive(self.left(idx), l, mid, a, b)
q2=self.query_recursive(self.right(idx), mid+1, r, a, b)
returnmax(q1, q2)

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


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.showData()