C++ Coding Test Course, Segment Tree

In recent algorithm coding tests, the segment tree is one of the frequently appearing data structures. The segment tree is a data structure that helps to efficiently calculate the sum, minimum, maximum, etc., of specific intervals. In this article, we will examine algorithm problems using segment trees and provide the process of solving them along with actual implementation examples.

Problem Description

Write a program that can perform queries to compute the sum of the interval [l, r] from a given array. The input includes the size of the array n and the number of queries m, followed by n integers representing the array. Each query consists of two integers l and r.

Example Problem

Input:
5 3
1 2 3 4 5
1 3
2 4
1 5

Output:
6
9
15

Problem Approach

The above problem can be solved by simply traversing the array to calculate the sum, but this is inefficient when the number of queries is large and the size of the array is large. Therefore, using a segment tree allows for an efficient solution to the problem.

What is a Segment Tree?

A segment tree is a tree structure that divides the array into segments and stores information about each segment. This tree has the following characteristics:

Each node represents an interval of the array.
Each leaf node stores an element of the array.
Each internal node stores the sum of its child nodes.

By utilizing a segment tree, the sum of an interval can be computed with O(log n) time complexity. This efficiency becomes more pronounced as the number of elements in the array increases.

Segment Tree Implementation

Now, let’s solve the problem using a segment tree. Below is the code to implement a segment tree in C++.


#include 
#include 
using namespace std;

class SegmentTree {
private:
    vector tree, arr;
    int n;

public:
    SegmentTree(const vector& input) {
        arr = input;
        n = arr.size();
        tree.resize(4 * n);
        build(0, 0, n - 1);
    }

    void build(int node, int start, int end) {
        if (start == end) {
            // Leaf node
            tree[node] = arr[start];
        } else {
            int mid = (start + end) / 2;
            // Recursively build the segment tree
            build(2 * node + 1, start, mid);
            build(2 * node + 2, mid + 1, end);
            // Internal node will have the sum of the two children
            tree[node] = tree[2 * node + 1] + tree[2 * node + 2];
        }
    }

    int sum(int l, int r) {
        return sum(0, 0, n - 1, l, r);
    }

    int sum(int node, int start, int end, int l, int r) {
        if (r < start || end < l) {
            // range represented by a node is completely outside the given range
            return 0;
        }
        if (l <= start && end <= r) {
            // range represented by a node is completely inside the given range
            return tree[node];
        }
        // range represented by a node is partially inside and partially outside the given range
        int mid = (start + end) / 2;
        int left_child = sum(2 * node + 1, start, mid, l, r);
        int right_child = sum(2 * node + 2, mid + 1, end, l, r);
        return left_child + right_child;
    }
};

int main() {
    int n, m;
    cin >> n >> m;
    vector arr(n);
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
    }
    SegmentTree segment_tree(arr);
    
    for (int i = 0; i < m; ++i) {
        int l, r;
        cin >> l >> r;
        // Adjusting for 0-based indexing
        cout << segment_tree.sum(l - 1, r - 1) << endl;
    }
    
    return 0;
}

Code Analysis

The above code implements the segment tree using the following procedures:

Class Definition: Defines the SegmentTree class and declares necessary member variables, which include the tree vector to store the segment tree, the original array arr, and the size of the array n.
Constructor: The constructor receives the input array, initializes the segment tree, and calls the build method to construct the tree.
Tree Building: The build method recursively builds the tree. Leaf nodes store elements of the original array, while internal nodes store the sum of child nodes.
Sum Calculation: The sum method calculates the sum of a given interval. This method checks whether a node falls within the given interval and returns an appropriate value or recursively calculates the sum.
Main Function: Takes input from the user regarding the size of the array and the number of queries, inputs the array, creates a SegmentTree object, and outputs the results for each query.

Conclusion

Segment trees are a powerful tool for quickly calculating the sum of specific intervals. By implementing a segment tree in C++, we enhanced the efficiency needed for solving algorithm problems. This data structure can be utilized not only for simple sums but also for various queries, so it is important to master it through more practice.

Now you have the foundational knowledge to implement and use segment trees in C++. Explore the power of this data structure by solving various problems.