I have written a program to sum up a value between a range using segment tree and this tree also supports single index updates. I am wondering if I have followed all the python rules for names and other pythonic nuances. Please let me know your comments.
Tree has below functions:
- Modify - helps in modifying the individual array elements.
Query - given a range this function will return the sum of all the elements in the range. Such as given [3, 7) will sum up index 3 to 6 excluding 7.
def next_power_2(n): n -= 1 n |= n >> 1 n |= n >> 2 n |= n >> 4 n |= n >> 8 n |= n >> 16 n += 1 return n class segment_tree(object): def __init__(self, n, data): self.total_elements = n self.n = next_power_2(n) #segment tree is always balanced binary tree self.tree = [0] * self.n + data + [0] * (self.n - n) def build(self): for i in range(self.n-1, 0, -1): self.tree[i] = self.tree[2*i] + self.tree[2*i+1] def modify(self, pos, val): if pos < 0 or pos >= self.total_elements: print("check your inputs") return -1 pos = pos + self.n self.tree[pos] = val while pos/2 > 0: self.tree[pos/2] = self.tree[pos] + self.tree[pos^1] pos = pos/2 def query(self, left, right): if left > self.total_elements or left < 0: print("check your inputs") return -1 if right > self.total_elements or right < 0: print("check your inputs") return -1 left = self.n + left right = self.n + right s = 0 while left < right: if left&1: s += self.tree[left] left += 1 if right&1: right -= 1 s += self.tree[right] left /= 2 right /= 2 return s def __repr__(self): return " ".join(str(i) for i in self.tree) print("enter the number of elements") n = int(raw_input()) print("enter the elements as space seperated integer values") data = [int(i) for i in raw_input().split()] st = segment_tree(n, data) st.build() print("enter the position and value which you want to modify as space seperated integer values") pos, val = map(int, raw_input().split()) print("pos", pos, "value", val) st.modify(pos, val) print("enter the left and right range as space seperated integer values") left, right = map(int, raw_input().split()) print("left", left, "right", right) print("sum for the given range is:", st.query(left, right)) print(st)