python coding test course, segment tree

Written on:

Table of Contents

  1. 1. Introduction to Segment Trees
  2. 2. Problem Description
  3. 3. Solution Process
  4. 4. Time Complexity Analysis
  5. 5. Conclusion

1. Introduction to Segment Trees

A segment tree is a data structure designed for efficient interval query processing and updates.
It is primarily designed for quickly handling sums, minimums, maximums, etc., over intervals in datasets like arrays.
The segment tree is structured as a complete binary tree, and each node of the tree stores information about a specific interval.

The main features of a segment tree are as follows:

  • Interval Query: It allows for quick retrieval of values within a specified range.
  • Update Functionality: The tree can be efficiently updated whenever the data changes.
  • Relatively Low Memory Usage: It requires relatively less memory compared to using an array.

2. Problem Description

Consider the following problem. Given an integer array arr, write a program that supports two functionalities: a query to find the sum of a specific interval [l, r] and an update to set the value at the ith position to val.

The input format for the problem is as follows:

  • First line: Size of the array N (1 ≤ N ≤ 100,000)
  • Second line: Elements of the array arr[1], arr[2], ..., arr[N]
  • Third line: Number of queries Q
  • Next Q lines: Three integers type, l, r for each query (type = 1: interval sum query, type = 2: update query)

For example, consider the following input:

5
1 2 3 4 5
3
1 1 3
2 2 10
1 1 5
        

Here, the first query requests the sum of the interval [1, 3], the second query updates the second element to 10, and the third query requests the sum of the interval [1, 5] after the update.

3. Solution Process

To solve this problem using a segment tree, the following approach can be used.

3.1. Constructing the Segment Tree

First, we need to construct the segment tree from the given array arr.
The parent node stores the sum of its child nodes’ values.
The tree can be initialized as follows:

class SegmentTree:
    def __init__(self, data):
        self.n = len(data)
        self.tree = [0] * (2 * self.n)
        # Insert data into leaf nodes
        for i in range(self.n):
            self.tree[self.n + i] = data[i]
        # Calculate internal nodes
        for i in range(self.n - 1, 0, -1):
            self.tree[i] = self.tree[i * 2] + self.tree[i * 2 + 1]
        

3.2. Processing Interval Sum Queries

To process interval sum queries, we need to traverse from the leaf nodes up to the root node.
To get the sum over the interval [l, r], we can implement as follows:

    def query(self, l, r):
        result = 0
        l += self.n
        r += self.n + 1
        while l < r:
            if l % 2 == 1:
                result += self.tree[l]
                l += 1
            if r % 2 == 1:
                r -= 1
                result += self.tree[r]
            l //= 2
            r //= 2
        return result
        

3.3. Processing Updates

An update query modifies the value at a specific index, affecting that node and its parent nodes.
It can be implemented as follows:

    def update(self, index, value):
        index += self.n
        self.tree[index] = value
        while index > 1:
            index //= 2
            self.tree[index] = self.tree[index * 2] + self.tree[index * 2 + 1]
        

3.4. Complete Code

Now let's write the complete code that includes all the components above:

def main():
    import sys
    input = sys.stdin.read
    data = input().split()
    
    idx = 0
    N = int(data[idx]); idx += 1
    arr = [0] * N
    for i in range(N):
        arr[i] = int(data[idx]); idx += 1
    Q = int(data[idx]); idx += 1
    
    seg_tree = SegmentTree(arr)
    
    output = []
    for _ in range(Q):
        query_type = int(data[idx]); idx += 1
        l = int(data[idx]); idx += 1
        r = int(data[idx]); idx += 1
        if query_type == 1:
            result = seg_tree.query(l - 1, r - 1)
            output.append(str(result))
        elif query_type == 2:
            seg_tree.update(l - 1, r)
    
    print('\n'.join(output))

if __name__ == "__main__":
    main()
        

4. Time Complexity Analysis

The time complexity of the segment tree is as follows:

  • Constructing the segment tree: O(N)
  • Interval sum query: O(log N)
  • Update query: O(log N)

Therefore, this algorithm can operate efficiently even with large datasets.

5. Conclusion

In this article, we explored how to handle interval sum queries and update queries using segment trees.
Segment trees are a powerful data structure that can be used effectively in various problems.
During coding tests, it is advisable to consider segment trees when facing interval query-related problems.