How to Implement Segment Tree in Python?

In this tutorial, we will learn what Segment Tree is and how to implement Segment Tree in Python with some non-recursive functions. This is a very important topic in Practical Data Structures.

Segment Tree is basically a data structure. It can be used to perform range queries and updates in an easy and fastest way.

Python program to implement segment tree

To understand Segment Tree we have to take an array first.

Let’s take an array A=[1,3,5,6,7,-3,6,2] of length 8 indexed from 0 to 7 and we have to solve problems called range queries and updates.

  1. Range queries mean to determine the sum of different segments of the given array.
    Example-sum(0,3)=1+3+5+6=15  (Here 0 and 3 represent the index no. of the given array).
  2. update means to change the value of a specified element of the given array to a new value.
    Example-If we perform an update(3,5), then the array becomes A=[1,3,5,5,7,-3,6,2] (Here 3 represents the index of the array, whose value to be changed and 5 represent the new or updated value).
  3. After performing the update, the sum(0,3) becomes 14(1+3+5+5) because of updating the value of the element index 3.

So, We can use Segment Tree to perform both the operations(range queries and update) in O(log n) time. First of all, let us have a look at the Segment Tree of the given array :

hand written draw of python array

In the above figure [L,R) indicates that left (L)is included and right(R)is excluded.

From the above image, you can see that the tree has 15 nodes in total and if you chose any one node from the parent nodes, let us suppose node 4 from the above tree then the left and right child of that node are node 8 and node 9 respectively.

So, in general, we can say that if we construct a Segment Tree for an element array total number of elements of the tree array will be (2*n-1) and the left and right child of the pth node will be at  2*p and 2*p+1 index respectively. Leaf nodes are starting from index (n) to (2*n-1). We can also observe that an element will be at index (k+n) in the segment tree array if it is kth element or k index element in the original array.

To perform range queries and updates using the Segment tree I am using three non-recursive functions. Python code of these three functions are given below:

# function to build the segmenttree array
def buildTree(a):
    # insert leaf nodes in tree
    for i in range(n):
        tree[n + i] = a[i]

    # creating parent node by adding left and right child
    for i in range(n - 1, 0, -1):
        tree[i] = tree[2*i] + tree[2*i+1]
# function to update a node of the tree
def updateTree(index, value):
    # set value at position index 
    tree[index + n] = value
    index+=n

    # after updating the child node,update parents
    i = index

    while i > 1: 
    #update parent by adding new left and right child
        tree[i//2] = tree[i] + tree[i+1]
        i =i//2

#function to find sum on different range 
def queryTree(l, r):
    sum = 0

    #to find the sum in the range [l,r)
    l += n
    r += n

    while l < r:

        if ((l & 1)>0):
            sum += tree[l]
            l += 1

        if ((r & 1)>0):
            r -= 1
            sum += tree[r]

        l =l// 2
        r =r// 2

    return sum

To check these three non-recursive function we have to write the main function.

if __name__ == "__main__":
    A = [1, 2, 3, 4, 5, 6, 7,8]
    n = len(A)
    buildTree(A)
    print(queryTree(1, 4))
    updateTree(2, 5)
    print(queryTree(1, 4))

Output:

9
11

Also, read: How to Add Trailing Zeros to String in Python?

 

Leave a Reply

Your email address will not be published. Required fields are marked *