Python/data_structures/binary_tree/segment_tree.py at master · imanghotbi-python/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
72
73
74
75
76
77
78
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): # noqa: E741
ifl==r: # noqa: E741
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): # noqa: E741
"""
 update(1, 1, N, a, b, v) for update val v to [a,b]
 """
ifr<aorl>b:
returnTrue
ifl==r: # noqa: E741
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): # noqa: E741
"""
 query(1, 1, N, a, b) for query max of [a,b]
 """
ifr<aorl>b:
return-math.inf
ifl>=aandr<=b: # noqa: E741
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)

defshow_data(self):
show_list= []
foriinrange(1, 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()