SegmentTree in python

segment_tree.py

#!/usr/bin/env python

# Inspired by: https://www.youtube.com/watch?v=Oq2E2yGadnU
# What i did was moving the first index to 0.

class SegmentTree(object):
    def __init__(self, data):
        n = len(data)
        self.offset = n - 1
        self.nodes = [0 for _ in range(n * 2 - 1)]
        for i in range(n):
            self.nodes[self.offset + i] = data[i]
        for i in range(self.offset-1, -1, -1):
            self.nodes[i] = self.nodes[i*2+1] + self.nodes[i*2+2]

    def get_sum(self, left, right):
        # [left, right)
        left += self.offset
        right += self.offset
        ret = 0
        while left < right:
            if left & 1 == 0: # on right child
                ret += self.nodes[left]
                left += 1
            if right & 1 == 0: # on right child
                right -= 1
                ret += self.nodes[right]
            left //= 2
            right //= 2
        return ret

import random
N = 100
nums = [i for i in range(N)]

for _ in range(1000):
    random.shuffle(nums)
    t = SegmentTree(nums)
    for i in range(N):
        for j in range(i+1, N):
            a = t.get_sum(i,j)
            b = sum(nums[i:j])
            assert a==b, '%d!=%d, %s %s %s'%(a, b, nums, i, j)