-
Notifications
You must be signed in to change notification settings - Fork 4
/
sumRangeSegmentTree.py
96 lines (80 loc) · 3.72 KB
/
sumRangeSegmentTree.py
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
from copy import deepcopy
from sys import maxint
from Tree.segment_tree.base_segment_tree import SegmentTree
class SumRangeSegmentTree(SegmentTree):
def __init__(self, input_arr):
super(SumRangeSegmentTree, self).__init__(input_arr)
# Time complexity: O(n)
# Space complexity: O(n)
def __build_tree_util__(self, low, high, pos):
if low == high:
self.seg_tree_arr[pos] = self.input_arr[low]
return
mid = (low + high)/2
self.__build_tree_util__(low, mid, 2*pos + 1)
self.__build_tree_util__(mid + 1, high, 2*pos + 2)
self.seg_tree_arr[pos] = self.seg_tree_arr[2*pos + 1] + self.seg_tree_arr[2*pos + 2]
def build_tree(self):
self.__build_tree_util__(0, len(self.input_arr)-1, 0)
return deepcopy(self.seg_tree_arr)
# Time complexity: O(log(n))
def __range_sum_query__(self, qs, qe, low, high, pos):
# Complete overlap, query range (qs, qe) is completely overlapping (low, high). e.g. (0, 4) (1, 3)
# return current value i.e. at index `pos`
if qs <= low and qe >= high:
return self.seg_tree_arr[pos]
# No overlap, query range (qs, qe) is not overlapping (low, high). e.g. (3, 4) (0, 2)
# Return 0
elif qs > high or qe < low:
return 0
# Partial overlap, query range (qs, qe) is partially overlapping (low, high). e.g. (3, 5) (1, 3)
# Go recursively in left & right subtree.
else:
mid = (low + high)/2
return (self.__range_sum_query__(qs, qe, low, mid, 2*pos + 1)
+ self.__range_sum_query__(qs, qe, mid+1, high, 2*pos + 2))
def range_sum_query(self, qs, qe):
range_sum = self.__range_sum_query__(qs, qe, 0, len(self.input_arr)-1, 0)
return range_sum
# Time complexity: O(log(n))
def __update_value__(self, i, diff, low, high, pos):
# If the input index `i` lies outside the range of this segment (low, high),
# return
if i < low or i > high:
return
# Leaf node: low and high are same,
# update it's value.
if low == high:
if low == i: # or high == i, as low & high are same
self.seg_tree_arr[pos] += diff
# If the input index `i` lies under the range of this segment (low, high), i.e.
# low <= i && i <= high && low != high
else:
mid = (low + high) / 2
self.__update_value__(i, diff, low, mid, 2 * pos + 1)
self.__update_value__(i, diff, mid + 1, high, 2 * pos + 2)
self.seg_tree_arr[pos] = self.seg_tree_arr[2 * pos + 1] + self.seg_tree_arr[2 * pos + 2]
def update_value(self, i, diff):
"""
:param i: Value/elt at index `i` needs to updated.
:param diff:
:return:
"""
self.__update_value__(i, diff, 0, len(self.input_arr) - 1, 0)
return deepcopy(self.seg_tree_arr)
if __name__ == '__main__':
input_arr = [1, 4, -1, 0]
seg_tree_obj = SumRangeSegmentTree(input_arr)
seg_tree_arr = seg_tree_obj.build_tree()
print "\n---------- Before modification ----------\n"
print "Sum-Range-Segment tree is:", seg_tree_arr
qs = 0; qe = 2
range_sum = seg_tree_obj.range_sum_query(qs, qe)
print "Sum of elements in range ({}, {}) is {}".format(qs, qe, range_sum)
modify_index = 2; diff = 3
print "\n---------- After modification (Adding {} at index {}) ----------\n".format(diff, modify_index)
seg_tree_arr_modified = seg_tree_obj.update_value(modify_index, diff)
print "Sum-Range-Segment tree is:", seg_tree_arr_modified
qs = 0; qe = 2
min_elt = seg_tree_obj.range_sum_query(qs, qe)
print "Sum of elements in range ({}, {}) is {}".format(qs, qe, min_elt)