Mastering Segment Trees in C++

Segment Trees are one of the most powerful and versatile data structures used in competitive programming and algorithmic problem solving. They allow us to answer range query problems efficiently while also supporting flexible update operations. In this comprehensive guide, we will dive deep into the mechanics of Segment Trees, their implementation in C++, and how to handle advanced techniques like Lazy Propagation.

What is a Segment Tree?

A Segment Tree is a binary tree used for storing information about intervals or segments. It allows querying which of the stored segments contain a given point. It is, in principle, a static structure; that is, it's a structure that cannot be modified once it's built. However, similar to other tree-based structures, we can modify it to support dynamic updates.

The most common application of Segment Trees is to solve the Range Minimum Query (RMQ) problem or Range Sum Query problem.

Why Segment Trees?

Imagine you have an array A of size N. You need to perform two types of operations:

1. Update: Change the value at index i to v.
2. Query: Find the sum (or minimum, maximum, etc.) of elements in the range [L, R].

If we use a simple array:

• Update: O(1)
• Query: O(N)

If we use a Prefix Sum array (for sum queries):

• Update: O(N) (need to rebuild prefix sums)
• Query: O(1)

A Segment Tree balances these operations:

• Update: O(log N)
• Query: O(log N)

This balance makes it ideal for scenarios with frequent updates and queries.

Structure of a Segment Tree

A Segment Tree is a binary tree where each node represents an interval.

• The root represents the entire array range [0, N-1].
• Each leaf represents a single element [i, i].
• The internal nodes represent the union of their children's intervals.

For an array of size N, the Segment Tree will have at most 4*N nodes.

Building the Segment Tree

To build the tree, we use a recursive approach. We start at the root and divide the range into two halves until we reach the leaf nodes.

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int MAXN = 100005;
int tree[4 * MAXN];
int arr[MAXN];

// Function to build the Segment Tree
// node: current node index in the tree array
// start: start index of the range covered by this node
// end: end index of the range covered by this node
void build(int node, int start, int end) {
    if (start == end) {
        // Leaf node will have a single element
        tree[node] = arr[start];
    } else {
        int mid = (start + end) / 2;
        // Recurse on the left child
        build(2 * node, start, mid);
        // Recurse on the right child
        build(2 * node + 1, mid + 1, end);
        // Internal node will have the sum of both of its children
        tree[node] = tree[2 * node] + tree[2 * node + 1];
    }
}

Point Updates

Updating a value in the array requires updating the leaf node corresponding to that index and then updating all its ancestors up to the root to reflect the change.

// Function to update a value in the Segment Tree
// idx: index of the element to be updated
// val: new value
void update(int node, int start, int end, int idx, int val) {
    if (start == end) {
        // Leaf node
        arr[idx] = val;
        tree[node] = val;
    } else {
        int mid = (start + end) / 2;
        if (start <= idx && idx <= mid) {
            // If idx is in the left child, recurse on the left child
            update(2 * node, start, mid, idx, val);
        } else {
            // If idx is in the right child, recurse on the right child
            update(2 * node + 1, mid + 1, end, idx, val);
        }
        // Update internal node
        tree[node] = tree[2 * node] + tree[2 * node + 1];
    }
}

Range Queries

To query a range [l, r], we traverse the tree. For each node covering a range [start, end]:

1. If the node's range is completely outside [l, r], return 0 (or identity element).
2. If the node's range is completely inside [l, r], return the node's value.
3. Otherwise, the range overlaps partially. Recurse on both children and combine the results.

// Function to query the Segment Tree
// l, r: query range
int query(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
    int mid = (start + end) / 2;
    int p1 = query(2 * node, start, mid, l, r);
    int p2 = query(2 * node + 1, mid + 1, end, l, r);
    return p1 + p2;
}

Lazy Propagation (Range Updates)

What if we need to update a range of values [l, r] by adding v to all of them? Doing this with point updates would take O(N log N) in the worst case. We can optimize this to O(log N) using Lazy Propagation.

The idea is to update a node only when needed. If we update a range, we mark the node as "lazy" indicating that there is a pending update for its children. We push this update down to the children only when we visit the node again.

Implementation with Lazy Propagation

We need an additional array lazy[] to store pending updates.

int lazy[4 * MAXN];

void push(int node, int start, int end) {
    if (lazy[node] != 0) {
        // Apply the pending update to the current node
        tree[node] += (end - start + 1) * lazy[node];
        
        // If it is not a leaf node, push the lazy value to its children
        if (start != end) {
            lazy[2 * node] += lazy[node];
            lazy[2 * node + 1] += lazy[node];
        }
        
        // Reset the lazy value of the current node
        lazy[node] = 0;
    }
}

void updateRange(int node, int start, int end, int l, int r, int val) {
    push(node, start, end); // Apply any pending updates
    
    if (start > end || start > r || end < l)
        return;
        
    if (start >= l && end <= r) {
        // Segment is fully within range
        tree[node] += (end - start + 1) * val;
        if (start != end) {
            lazy[2 * node] += val;
            lazy[2 * node + 1] += val;
        }
        return;
    }
    
    int mid = (start + end) / 2;
    updateRange(2 * node, start, mid, l, r, val);
    updateRange(2 * node + 1, mid + 1, end, l, r, val);
    tree[node] = tree[2 * node] + tree[2 * node + 1];
}

int queryRange(int node, int start, int end, int l, int r) {
    push(node, start, end); // Apply any pending updates
    
    if (start > end || start > r || end < l)
        return 0;
        
    if (start >= l && end <= r)
        return tree[node];
        
    int mid = (start + end) / 2;
    int p1 = queryRange(2 * node, start, mid, l, r);
    int p2 = queryRange(2 * node + 1, mid + 1, end, l, r);
    return p1 + p2;
}

Complexity Analysis

• Space Complexity: O(N) (specifically 4*N nodes).
• Time Complexity:
  • Build: O(N)
  • Point Update: O(log N)
  • Range Query: O(log N)
  • Range Update (Lazy): O(log N)

Advanced Applications

Segment Trees are extremely flexible. Beyond simple sum or min/max queries, they can be used for:

1. Finding the K-th zero in an array.
2. Merging Sort Trees (storing a sorted vector in each node) to answer queries like "count numbers greater than K in range [L, R]".
3. Persistent Segment Trees to query past versions of the array.
4. 2D Segment Trees for matrix queries.

Conclusion

Segment Trees are a fundamental tool in a competitive programmer's arsenal. While they may seem complex at first, understanding the recursive structure and the logic behind lazy propagation unlocks the ability to solve a wide class of efficient range query problems.

Mastering this data structure requires practice. Try implementing it from scratch and solving problems on platforms like Codeforces, LeetCode, or SPOJ to solidify your understanding.

Happy Coding!