Demystifying Segment Trees: An Efficient Solution for Range Queries

Introduction

Data lies at the heart of any application and as software engineers we are constantly challenged to optimize data manipulation and efficiently extract valuable insights. This is where Segment Trees come into play.

A Segment Tree is a data structure designed to handle range queries and updates on dynamic sets. At its core, it breaks down an array into smaller segments, allowing us to perform aggregate operations within specified intervals efficiently.

They play a pivotal role in a number of real-world scenarios. From database management to image processing and GIS applications, the ability to quickly compute sums, minimums, maximums, and more within specific ranges makes Segment Trees a critical asset for optimizing performance and resource management.

Building a Segment Tree

Let’s say we are given an array and we want to support two features — update & query:

Array: [5, 3, 7, 1, 4, 2]

Update: update(index, val)

Query: querySumRange(leftIndex, rightIndex)

Naive Approach: We could update the array in constant time and query naively by iterating through every element in the range and calculating the total sum which would be a O(n) operation. And as n gets larger, this may not be efficient.

Better Approach: Lets create a tree based on ranges in the array and update & query using this tree.

Here each node represents the left and right indices of the subarray. The root node is the entire array, the left child is the left half of the root node’s indices and the right child is the right half of the root node’s indices.

The main goal of this tree is to efficiently query sums based on an index range, so naturally the base case (leaf node) would be if the subarray size is 1. So let’s update our tree.

Now each node contains the start and end indices and also the sum over that range. Here’s an example of how we could implement this.

class Node {
  constructor(sum, l, r) {
    this.sum = sum; // sum over the range [l, r]
    this.leftIdx = l; // left index
    this.rightIdx = r; // right index
    this.left = null; // left node
    this.right = null; // right node
  }
}

class SegmentTree {
  constructor(array) {
    this.segmentTree = this.buildTree(array, 0, array.length - 1);
  }
  // O(n) time
  buildTree(array, l, r) {
    // base case if l equals r which denotes a subarray length of 1
    if (l === r) return new Node(array[l], l, r);
    // get mid point
    const midIdx = Math.floor((l + r) / 2);
    // create root node with an initial sum of 0
    const root = new Node(0, l, r);
    // assign root.left to recursively calling buildTree on left half of array
    root.left = this.buildTree(array, l, midIdx);
    // assign root.right to recursively calling buildTree on right half of array
    root.right = this.buildTree(array, midIdx + 1, r);
    // update sum of root
    root.sum = root.left.sum + root.right.sum;
    return root;
  }
}

In our implementation, build tree will take O(n) time.

Updating a Segment Tree

Let’s say we want to update index 3 to be the value 4. In our tree we would want to go to the leaf node that represents index 3. We could use a form of binary search to do this efficiently in O(log(n)) time.

Keep in mind that once we update the leaf node, we would need to propagate that new value all the way back up the tree (i.e update the sums of the appropriate parent nodes).

The implementation would look something like the following:

// O(log(n)) time
update(index, val, current = this.segmentTree) {
    // base case if current left and right indices are equal, update sum
    if (current.leftIdx === current.rightIdx) {
      current.sum = val;
      return;
    }
    // calculate mid point
    const midIdx = Math.floor((current.leftIdx + current.rightIdx) / 2);
    // index > midIdx, explore right subtree. Otherwise, explore left subtree
    if (index > midIdx) {
      this.update(index, val, current.right);
    } else {
      this.update(index, val, current.left);
    }
    // update the currentNode's sum
    current.sum = current.left.sum + current.right.sum;
  }

Querying a Segment Tree

Now let’s say we want to get the sum of the range from index 2 to 4. We can do so in O(log(n)) time by traversing the tree until we hit the entire range we are querying.

Here’s one way to implement this:

// O(log(n)) time
rangeQuery(leftIndex, rightIndex, current = this.segmentTree) {
    // base case if current leftIdx = leftIndex && current rightIdx = rightIndex
    if (current.leftIdx === leftIndex && current.rightIdx === rightIndex)
      return current.sum;
    // calculate mid point
    const midIdx = Math.floor((current.leftIdx + current.rightIdx) / 2);
    // if leftIndex > midIdx, the entire range is in the right subtree
    if (leftIndex > midIdx) {
      return this.rangeQuery(leftIndex, rightIndex, current.right);
    } else if (rightIndex <= midIdx) {
      // if right index <= midIdx, the entire range is in left subtree
      return this.rangeQuery(leftIndex, rightIndex, current.left);
    } else {
      // otherwise the range spans both subtrees
      return (
        this.rangeQuery(leftIndex, midIdx, current.left) +
        this.rangeQuery(midIdx + 1, rightIndex, current.right)
      );
    }
  }

Conclusion

Hopefully by now you can see the power of Segment Trees and how to implement them. Not only is this data structure a great tool to have in your toolbox, but also incredibly applicable in every day life.

Thanks for reading and feel free to follow or connect with me on Linkedin and Github. If you have any suggestions on what you would like me to write about next, feel free to leave a comment!