-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathsegment_tree.py
More file actions
71 lines (65 loc) · 1.84 KB
/
segment_tree.py
File metadata and controls
71 lines (65 loc) · 1.84 KB
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
class SegmentTree(object):
def __init__(self, arr):
self.lst = arr
self.s = [0 for _ in range(4*len(self.lst))] # segment
def build(self, i, l, r):
"""
Builds the segment tree
:param i: Index of segment tree
:param l: range start
:param r: range end
:return: None
"""
if l == r:
self.s[i] = self.lst[l]
return
mid = l + (r-l)//2
self.build(i*2+1, l, mid)
self.build(i*2+2, mid+1, r)
self.s[i] = self.s[i*2+1] + self.s[i*2+2]
return
def modify(self, p, x, i, l, r):
"""
Modify the segment tree
:param p: position in array, modified with
:param x: value x
:param i: index of segment tree
:param l: range start
:param r: range end
:return: None
"""
self.s[i] += x - self.lst[p]
if l == r:
self.s[i] = x
self.lst[p] = x
return
mid = l + (r-l)//2
if p < mid:
self.modify(p, x, i*2+1, l, mid)
else:
self.modify(p, x, i*2+2, mid+1, r)
return
def query(self, x, y, i, l, r):
"""
Query the result of range
:param x: search range start
:param y: search range end
:param i: index of segment
:param l: range start
:param r: range end
:return: None
"""
if x > r or l > y:
return 0
elif x <= l and r <= y:
return self.s[i]
else:
mid = l + (r-l)//2
return self.query(x, y, i*2+1, l, mid) + self.query(x, y, i*2+2, mid+1, r)
if __name__ == "__main__":
S = SegmentTree([1, 2, 3, 4])
S.build(0, 0, 3)
print(S.s)
S.modify(0, 4, 0, 0, 3)
print(S.lst, S.s)
print (S.query(0, 2, 0, 0, 3))