A simple approach to segment trees

A segment tree is a tree data structure that allows aggregation queries and updates over array intervals in logarithmic time. As I see it, there are three major use cases for segment trees:

1. Static segment trees: This is probably the most common use case. We preprocess an array of N elements to construct a segment tree in O(N). Now, we can query aggregates over any arbitrary range/segment of the array in O(log N).
2. Segment tree with point updates: This allows us to update array values, one at a time in O(log N), while still maintaining the segment tree structure. Queries over any arbitrary range still occurs in O(log N).
3. Segment tree with range updates: This allows us to update a range of array elements at once in O(N) in the worst case, however problem specific optimizations and lazy propagation typically give huge improvements. Queries over any arbitrary range still occurs in O(log N).

In this post, I’ll cover the first two use cases because they go together. Given a static segment tree, it is very easy to add point update capability to it. I’ll leave the third use case as the subject matter of a future blog post. I intend this post to be a practical introduction to segment trees, rather than a theoretical description, so it will focus on how we can divide a segment tree into its components, the working of each component and how we can separate the problem specific logic from the underlying data structure. We’ll build a template for a segment tree and then apply it to several problems to understand how problem specific logic can be cleanly separated from the template.

Structure of a segment tree
Let’s understand what a segment tree looks like. Each node in a segment tree stores aggregate statistics for some range/segment of an array. The leaf nodes stores aggregate statistics for individual array elements. Although a segment tree is a tree, it is stored in an array similar to a heap. If the input array had 2ⁿ elements (i.e., the number of elements were a power of 2), then the segment tree over it would look something like this:

Each node here shows the segment of the input array for which it is responsible. The number outside a node indicates its index in the segment tree array. Clearly, if the array size N were a power of 2, then the segment tree would have 2*N-1 nodes. It is simpler to store the first node at index 1 in the segment tree array in order to simplify the process of finding indices of left and right children (a node at index i has left and right children at 2*i and 2*i+1 respectively). Thus, for an input array of size N, an array of size 2*N would be required to store the segment tree.

In practice, however, N is not usually a power of 2, so we have to find the power of 2 immediately greater than N, let’s call it x, and allocate an array of size 2*x to store the segment tree. The following procedure calculates the size of array required to store a segment tree for an input array size N:

	int getSegmentTreeSize(int N) {
	int size = 1;
	for (; size < N; size <<= 1);
	return size << 1;
	}

view raw Size of a segment tree.cpp hosted with ❤ by GitHub

We’ll try to separate the implementation of the underlying data structure from the problem specific logic. For this purpose, let us define a structure for a segment tree node:

	struct SegmentTreeNode {
	// variables to store aggregate statistics and
	// any other information required to merge these
	// aggregate statistics to form parent nodes
	void assignLeaf(T value) {
	// T is the type of input array element
	// Given the value of an input array element,
	// build aggregate statistics for this leaf node
	}
	void merge(SegmentTreeNode& left, SegmentTreeNode& right) {
	// merge the aggregate statistics of left and right
	// children to form the aggregate statistics of
	// their parent node
	}
	V getValue() {
	// V is the type of the required aggregate statistic
	// return the value of required aggregate statistic
	// associated with this node
	}
	};

view raw Segment tree node.cpp hosted with ❤ by GitHub

Building a segment tree
We can build a segment tree recursively in a depth first manner, starting at the root node (representative of the whole input array), working our way towards the leaves (representatives of individual input array elements). Once both children of a node have returned, we can merge their aggregate statistics to form their parent node.

	void buildTree(T arr[], int stIndex, int lo, int hi) {
	if (lo == hi) {
	nodes[stIndex].assignLeaf(arr[lo]);
	return;
	}
	int left = 2 * stIndex, right = left + 1, mid = (lo + hi) / 2;
	buildTree(arr, left, lo, mid);
	buildTree(arr, right, mid + 1, hi);
	nodes[stIndex].merge(nodes[left], nodes[right]);
	}

view raw Building a segment tree.cpp hosted with ❤ by GitHub

Here I’ve assumed that the type of input array elements is T. stIndex represents the index of current segment tree node in the segment tree array, lo and hi indicate the range/segment of input array this node is responsible for. We build the whole segment tree with a single call to buildTree(arr, 1, 0, N-1), where N is the size of input array arr. Clearly, the time complexity of this procedure is O(N), assuming that assignLeaf() and merge() operations work in O(1).

Querying the segment tree
Suppose we want to query the aggregate statistic associated with the segment [lo,hi], we can do this recursively as follows:

	// V is the type of the required aggregate statistic
	V getValue(int lo, int hi) {
	SegmentTreeNode result = getValue(1, 0, N-1, lo, hi);
	return result.getValue();
	}
	// nodes[stIndex] is responsible for the segment [left, right]
	// and we want to query for the segment [lo, hi]
	SegmentTreeNode getValue(int stIndex, int left, int right, int lo, int hi) {
	if (left == lo && right == hi)
	return nodes[stIndex];
	int mid = (left + right) / 2;
	if (lo > mid)
	return getValue(2*stIndex+1, mid+1, right, lo, hi);
	if (hi <= mid)
	return getValue(2*stIndex, left, mid, lo, hi);
	SegmentTreeNode leftResult = getValue(2*stIndex, left, mid, lo, mid);
	SegmentTreeNode rightResult = getValue(2*stIndex+1, mid+1, right, mid+1, hi);
	SegmentTreeNode result;
	result.merge(leftResult, rightResult);
	return result;
	}

view raw Querying the segment tree.cpp hosted with ❤ by GitHub

This procedure is similar to the one used for building the segment tree, except that we cut off recursion when we reach a desired segment. The complexity of this procedure is O(log N).

Updating the segment tree
The above two procedures, building the segment tree and querying it, are sufficient for the first use case: a static segment tree. It so happens that the second use case: point updates, doesn’t require many changes. In fact, we don’t have to change the problem specific logic at all. No changes in the structure SegmentTreeNode are required.
We just need to add in a procedure for updating the segment tree. It is very similar to the buildTree() procedure, the only difference being that it follows only one path down the tree (the one that leads to the leaf node being updated) and comes back up, recursively updating parent nodes along this same path.

	// We want to update the value associated with index in the input array
	void update(int index, T value) {
	update(1, 0, N-1, index, value);
	}
	// nodes[stIndex] is responsible for segment [lo, hi]
	void update(int stIndex, int lo, int hi, int index, T value) {
	if (lo == hi) {
	nodes[stIndex].assignLeaf(value);
	return;
	}
	int left = 2 * stIndex, right = left + 1, mid = (lo + hi) / 2;
	if (index <= mid)
	update(left, lo, mid, index, value);
	else
	update(right, mid+1, hi, index, value);
	nodes[stIndex].merge(nodes[left], nodes[right]);
	}

view raw Updating the segment tree.cpp hosted with ❤ by GitHub

Clearly, the complexity of this operation is O(log N), assuming that assignLeaf() and merge() work in O(1).

Segment Tree template
Let’s put all this together to complete the template for a segment tree.

	// T is the type of input array elements
	// V is the type of required aggregate statistic
	template<class T, class V>
	class SegmentTree {
	SegmentTreeNode* nodes;
	int N;
	public:
	SegmentTree(T arr[], int N) {
	this->N = N;
	nodes = new SegmentTreeNode[getSegmentTreeSize(N)];
	buildTree(arr, 1, 0, N-1);
	}
	~SegmentTree() {
	delete[] nodes;
	}
	V getValue(int lo, int hi) {
	SegmentTreeNode result = getValue(1, 0, N-1, lo, hi);
	return result.getValue();
	}
	void update(int index, T value) {
	update(1, 0, N-1, index, value);
	}
	private:
	void buildTree(T arr[], int stIndex, int lo, int hi) {
	if (lo == hi) {
	nodes[stIndex].assignLeaf(arr[lo]);
	return;
	}
	int left = 2 * stIndex, right = left + 1, mid = (lo + hi) / 2;
	buildTree(arr, left, lo, mid);
	buildTree(arr, right, mid + 1, hi);
	nodes[stIndex].merge(nodes[left], nodes[right]);
	}
	SegmentTreeNode getValue(int stIndex, int left, int right, int lo, int hi) {
	if (left == lo && right == hi)
	return nodes[stIndex];
	int mid = (left + right) / 2;
	if (lo > mid)
	return getValue(2*stIndex+1, mid+1, right, lo, hi);
	if (hi <= mid)
	return getValue(2*stIndex, left, mid, lo, hi);
	SegmentTreeNode leftResult = getValue(2*stIndex, left, mid, lo, mid);
	SegmentTreeNode rightResult = getValue(2*stIndex+1, mid+1, right, mid+1, hi);
	SegmentTreeNode result;
	result.merge(leftResult, rightResult);
	return result;
	}
	int getSegmentTreeSize(int N) {
	int size = 1;
	for (; size < N; size <<= 1);
	return size << 1;
	}
	void update(int stIndex, int lo, int hi, int index, T value) {
	if (lo == hi) {
	nodes[stIndex].assignLeaf(value);
	return;
	}
	int left = 2 * stIndex, right = left + 1, mid = (lo + hi) / 2;
	if (index <= mid)
	update(left, lo, mid, index, value);
	else
	update(right, mid+1, hi, index, value);
	nodes[stIndex].merge(nodes[left], nodes[right]);
	}
	};

view raw Segment tree template.cpp hosted with ❤ by GitHub

We shall now see how this template can be used to solve different problems, without requiring a change in the tree implementation, and how the structure SegmentTreeNode is implemented differently for different problems.

The first problem we’ll look at it is GSS1. This problem asks for a solution to maximum subarray problem for each range of an array. My objective here is not to explain how to solve this problem, rather to demonstrate how easily it can be implemented with the above template at hand.
As it turns out, we need to store 4 values in each segment tree node to be able to merge child nodes to form a solution to their parent’s node:

1. Maximum sum of a subarray, starting at the leftmost index of this range
2. Maximum sum of a subarray, ending at the rightmost index of this range
3. Maximum sum of any subarray in this range
4. Sum of all elements in this range

The SegmentTreeNode for this problem looks as follows:

	struct SegmentTreeNode {
	int prefixMaxSum, suffixMaxSum, maxSum, sum;
	void assignLeaf(int value) {
	prefixMaxSum = suffixMaxSum = maxSum = sum = value;
	}
	void merge(SegmentTreeNode& left, SegmentTreeNode& right) {
	sum = left.sum + right.sum;
	prefixMaxSum = max(left.prefixMaxSum, left.sum + right.prefixMaxSum);
	suffixMaxSum = max(right.suffixMaxSum, right.sum + left.suffixMaxSum);
	maxSum = max(prefixMaxSum, max(suffixMaxSum, max(left.maxSum, max(right.maxSum, left.suffixMaxSum + right.prefixMaxSum))));
	}
	int getValue() {
	return maxSum;
	}
	};

view raw GSS1 segment tree node.cpp hosted with ❤ by GitHub

The complete solution for this problem can be viewed here.

The second problem we’ll look at is GSS3, which is very similar to GSS1 with the only difference being that it also asks for updates to array elements, while still maintaining the structure for getting maximum subarray sum. Now, we can understand the advantage of separating problem specific logic from the segment tree implementation. This problem requires no changes to the template and even uses the same SegmentTreeNode as used for GSS1. The complete solution for this problem can be viewed here.

The third problem: BRCKTS, we’ll look at is very different from the first two but the differences are only superficial since we’ll be able to solve it using the same structure. This problem gives a string containing parenthesis (open and closed), requires making updates to individual parenthesis (changing an open parenthesis to closed or vice versa), and checking if the whole string represents a correct parenthesization.
As it turns out, we need only 2 things in each segment tree node:

1. The number of unmatched open parenthesis in this range
2. The number of unmatched closed parenthesis in this range

The SegmentTreeNode for this problem looks as follows:

	struct SegmentTreeNode {
	int unmatchedOpenParans, unmatchedClosedParans;
	void assignLeaf(char paranthesis) {
	if (paranthesis == '(')
	unmatchedOpenParans = 1, unmatchedClosedParans = 0;
	else
	unmatchedOpenParans = 0, unmatchedClosedParans = 1;
	}
	void merge(SegmentTreeNode& left, SegmentTreeNode& right) {
	int newMatches = min(left.unmatchedOpenParans, right.unmatchedClosedParans);
	unmatchedOpenParans = right.unmatchedOpenParans + left.unmatchedOpenParans - newMatches;
	unmatchedClosedParans = left.unmatchedClosedParans + right.unmatchedClosedParans - newMatches;
	}
	bool getValue() {
	return unmatchedOpenParans == 0 && unmatchedClosedParans == 0;
	}
	};

view raw BRCKTS segment tree node.cpp hosted with ❤ by GitHub

The complete solution for this problem can be viewed here.

The final problem we’ll look at in this post is KGSS. This problem asks for the maximum pair sum in each subarray and also requires updates to individual array elements. As it turns out, we only need to store 2 things in each segment tree node:

1. The maximum value in this range
2. The second maximum value in this range

The SegmentTreeNode for this problem looks as follows:

	struct SegmentTreeNode {
	int maxNum, secondMaxNum;
	void assignLeaf(int num) {
	maxNum = num;
	secondMaxNum = -1;
	}
	void merge(SegmentTreeNode& left, SegmentTreeNode& right) {
	maxNum = max(left.maxNum, right.maxNum);
	secondMaxNum = min(max(left.maxNum, right.secondMaxNum), max(right.maxNum, left.secondMaxNum));
	}
	int getValue() {
	return maxNum + secondMaxNum;
	}
	};

view raw KGSS segment tree node.cpp hosted with ❤ by GitHub

The complete solution for this problem can be viewed here.

I hope this post presented a gentle introduction to segment trees and I look forward to feedback for possible improvements and suggestions for a future post on segment trees with lazy propagation.

Continue to Part 2…