Segment Tree Tutorial

Introduction

A Segment Tree is a data structure used to efficiently perform range queries and updates on an array. It is particularly useful when you need to query the sum, minimum, maximum, or any associative operation over a range of elements in an array, and it supports these operations in O(log n) time.

Key Operations

1. Build: Construct the Segment Tree from an array.
2. Update: Update a specific element in the array and reflect this change in the Segment Tree.
3. Query: Query a specific range of the array for a sum, minimum, maximum, or any other associative function.

Segment Tree Structure

A Segment Tree is typically represented as a binary tree, where each node represents a segment (or range) of the array. The leaf nodes represent individual elements, and the internal nodes represent the result of the function (like sum, min, max) applied to the segment.

Implementation

Here’s a basic implementation of a Segment Tree in JavaScript for range sum queries:

class SegmentTree {
  constructor(arr) {
    this.n = arr.length;
    this.tree = new Array(2 * this.n);
    this.build(arr);
  }

  // Build the Segment Tree
  build(arr) {
    // Insert leaf nodes in the tree
    for (let i = 0; i < this.n; i++) {
      this.tree[this.n + i] = arr[i];
    }

    // Build the tree by calculating parents
    for (let i = this.n - 1; i > 0; i--) {
      this.tree[i] = this.tree[i * 2] + this.tree[i * 2 + 1];
    }
  }

  // Update the value at index i to newValue
  update(index, newValue) {
    index += this.n;
    this.tree[index] = newValue;

    // Recalculate the values in the tree
    while (index > 1) {
      index = Math.floor(index / 2);
      this.tree[index] = this.tree[2 * index] + this.tree[2 * index + 1];
    }
  }

  // Query the sum of values in range [left, right)
  query(left, right) {
    left += this.n; // add +1 to make the left query exclusive
    right += this.n; // add +1 to make the query right inclusive
    let sum = 0;

    while (left < right) {
      if (left % 2 === 1) {
        sum += this.tree[left];
        left++;
      }
      if (right % 2 === 1) {
        right--;
        sum += this.tree[right];
      }
      left = Math.floor(left / 2);
      right = Math.floor(right / 2);
    }

    return sum;
  }
}

// Example Usage
const arr = [1, 3, 5, 7, 9, 11];
const segTree = new SegmentTree(arr);

// Querying the sum from index 1 to 4 (3 + 5 + 7)
console.log(segTree.query(1, 4)); // Output: 15

// Updating the value at index 2 to 6
segTree.update(2, 6);

// Querying the sum from index 1 to 4 again (3 + 6 + 7)
console.log(segTree.query(1, 4)); // Output: 16

Using Bit Manipulation

class SegmentTree {
  constructor(arr) {
    this.n = arr.length;
    this.tree = new Array(2 * this.n);
    this.build(arr);
  }

  // Build the Segment Tree
  build(arr) {
    // Insert leaf nodes in the tree
    for (let i = 0; i < this.n; i++) {
      this.tree[this.n + i] = arr[i];
    }

    // Build the tree by calculating parents
    for (let i = this.n - 1; i > 0; i--) {
      this.tree[i] = this.tree[i << 1] + this.tree[i << 1 | 1];
    }
  }

  // Update the value at index i to newValue
  update(index, newValue) {
    index += this.n;
    this.tree[index] = newValue;

    // Recalculate the values in the tree
    while (index > 1) {
      index >>= 1; // Move up by halving index
      this.tree[index] = this.tree[index << 1] + this.tree[index << 1 | 1];
    }
  }

  // Query the sum of values in range [left, right)
  query(left, right) {
    left += this.n;
    right += this.n;
    let sum = 0;

    while (left < right) {
      if (left & 1) { // If left is odd, it is a right child
        sum += this.tree[left++];
      }
      if (right & 1) { // If right is odd, decrement it to include the left sibling
        sum += this.tree[--right];
      }
      left >>= 1; // Move up by halving index
      right >>= 1;
    }

    return sum;
  }
}

// Example Usage
const arr = [1, 3, 5, 7, 9, 11];
const segTree = new SegmentTree(arr);

// Querying the sum from index 1 to 4 (3 + 5 + 7)
console.log(segTree.query(1, 4)); // Output: 15

// Updating the value at index 2 to 6
segTree.update(2, 6);

// Querying the sum from index 1 to 4 again (3 + 6 + 7)
console.log(segTree.query(1, 4)); // Output: 16

Segment Tree implementation for both minimum and maximum

class SegmentTree {
  constructor(arr) {
    this.n = arr.length;
    // Size 4n to ensure enough space for both min and max trees
    this.minTree = new Array(4 * this.n).fill(Infinity);
    this.maxTree = new Array(4 * this.n).fill(-Infinity);
    this.buildTree(arr, 0, this.n - 1, 1);
  }

  // Build both min and max trees simultaneously
  buildTree(arr, start, end, node) {
    if (start === end) {
      this.minTree[node] = arr[start];
      this.maxTree[node] = arr[start];
      return;
    }

    const mid = Math.floor((start + end) / 2);
    // Build left subtree
    this.buildTree(arr, start, mid, 2 * node);
    // Build right subtree
    this.buildTree(arr, mid + 1, end, 2 * node + 1);

    // Update current node
    this.minTree[node] = Math.min(this.minTree[2 * node], this.minTree[2 * node + 1]);
    this.maxTree[node] = Math.max(this.maxTree[2 * node], this.maxTree[2 * node + 1]);
  }

  // Update value at index
  update(index, value) {
    this.#update(index, value, 0, this.n - 1, 1);
  }

  #update(index, value, start, end, node) {
    // Found the leaf node to update
    if (start === end) {
      this.minTree[node] = value;
      this.maxTree[node] = value;
      return;
    }

    const mid = Math.floor((start + end) / 2);
    if (index <= mid) {
      this.#update(index, value, start, mid, 2 * node);
    } else {
      this.#update(index, value, mid + 1, end, 2 * node + 1);
    }

    // Update the current node after child updates
    this.minTree[node] = Math.min(this.minTree[2 * node], this.minTree[2 * node + 1]);
    this.maxTree[node] = Math.max(this.maxTree[2 * node], this.maxTree[2 * node + 1]);
  }

  // Query minimum in range [left, right]
  queryMin(left, right) {
    return this.#queryMin(left, right, 0, this.n - 1, 1);
  }

  #queryMin(left, right, start, end, node) {
    // No overlap
    if (right < start || left > end) return Infinity;

    // Complete overlap
    if (left <= start && right >= end) return this.minTree[node];

    // Partial overlap - split the query
    const mid = Math.floor((start + end) / 2);
    const leftMin = this.#queryMin(left, right, start, mid, 2 * node);
    const rightMin = this.#queryMin(left, right, mid + 1, end, 2 * node + 1);
    return Math.min(leftMin, rightMin);
  }

  // Query maximum in range [left, right]
  queryMax(left, right) {
    return this.#queryMax(left, right, 0, this.n - 1, 1);
  }

  #queryMax(left, right, start, end, node) {
    // No overlap
    if (right < start || left > end) return -Infinity;

    // Complete overlap
    if (left <= start && right >= end) return this.maxTree[node];

    // Partial overlap - split the query
    const mid = Math.floor((start + end) / 2);
    const leftMax = this.#queryMax(left, right, start, mid, 2 * node);
    const rightMax = this.#queryMax(left, right, mid + 1, end, 2 * node + 1);
    return Math.max(leftMax, rightMax);
  }
}

// Example usage:
const arr = [1, 3, 2, 7, 9, 11, 8, 4];
const segTree = new SegmentTree(arr);

// Test cases
console.log("Min in range [0, 3]:", segTree.queryMin(0, 3));  // Should output 1
console.log("Max in range [0, 3]:", segTree.queryMax(0, 3));  // Should output 7
console.log("Min in range [4, 7]:", segTree.queryMin(4, 7));  // Should output 4
console.log("Max in range [4, 7]:", segTree.queryMax(4, 7));  // Should output 11

// Update value and test again
segTree.update(2, 10);
console.log("Min in range [0, 3] after update:", segTree.queryMin(0, 3));  // Should output 1
console.log("Max in range [0, 3] after update:", segTree.queryMax(0, 3));  // Should output 10

Segment Tree implementation with Lazy Propagation for Range Updates and Point Sum Tracking

class LazySegmentTree {
  constructor(arr) {
    this.n = arr.length;
    this.tree = new Array(4 * this.n).fill(0);
    this.lazy = new Array(4 * this.n).fill(0);
    this.buildTree(arr, 0, this.n - 1, 1);
  }

  // Build the initial segment tree
  buildTree(arr, start, end, node) {
    if (start === end) {
      this.tree[node] = arr[start];
      return;
    }

    const mid = Math.floor((start + end) / 2);
    this.buildTree(arr, start, mid, 2 * node);
    this.buildTree(arr, mid + 1, end, 2 * node + 1);
    this.tree[node] = this.tree[2 * node] + this.tree[2 * node + 1];
  }

  // Push lazy updates down to children
  pushDown(node, start, end) {
    if (this.lazy[node] !== 0) {
      this.tree[node] += (end - start + 1) * this.lazy[node];

      if (start !== end) {
        this.lazy[2 * node] += this.lazy[node];
        this.lazy[2 * node + 1] += this.lazy[node];
      }

      this.lazy[node] = 0;
    }
  }

  // Range update: Add 'value' to all elements in range [left, right]
  rangeUpdate(left, right, value) {
    this.#rangeUpdate(left, right, value, 0, this.n - 1, 1);
  }

  #rangeUpdate(left, right, value, start, end, node) {
    this.pushDown(node, start, end);

    // No overlap
    if (end < left || right < start) return;

    // Complete overlap
    if (left <= start && end <= right) {
      this.lazy[node] += value;
      this.pushDown(node, start, end);
      return;
    }

    // Partial overlap
    const mid = Math.floor((start + end) / 2);
    this.#rangeUpdate(left, right, value, start, mid, 2 * node);
    this.#rangeUpdate(left, right, value, mid + 1, end, 2 * node + 1);
    this.tree[node] = this.tree[2 * node] + this.tree[2 * node + 1];
  }

  // Point update: Set value at index
  pointUpdate(index, value) {
    this.#pointUpdate(index, value, 0, this.n - 1, 1);
  }

  #pointUpdate(index, value, start, end, node) {
    this.pushDown(node, start, end);

    if (start === end) {
      this.tree[node] = value;
      return;
    }

    const mid = Math.floor((start + end) / 2);
    if (index <= mid) {
      this.#pointUpdate(index, value, start, mid, 2 * node);
    } else {
      this.#pointUpdate(index, value, mid + 1, end, 2 * node + 1);
    }

    this.tree[node] = this.tree[2 * node] + this.tree[2 * node + 1];
  }

  // Query sum in range [left, right]
  querySum(left, right) {
    return this.#querySum(left, right, 0, this.n - 1, 1);
  }

  #querySum(left, right, start, end, node) {
    if (end < left || right < start) return 0;

    this.pushDown(node, start, end);

    if (left <= start && end <= right) return this.tree[node];

    const mid = Math.floor((start + end) / 2);
    const leftSum = this.#querySum(left, right, start, mid, 2 * node);
    const rightSum = this.#querySum(left, right, mid + 1, end, 2 * node + 1);
    return leftSum + rightSum;
  }

  // Get value at a specific index
  getValue(index) {
    return this.querySum(index, index);
  }
}

// Example usage and test cases
const arr = [1, 3, 5, 7, 9, 11];
const segTree = new LazySegmentTree(arr);

// Test 1: Basic sum query
console.log("Initial sum [0, 2]:", segTree.querySum(0, 2));  // Should be 9 (1 + 3 + 5)

// Test 2: Range update
segTree.rangeUpdate(1, 3, 2);  // Add 2 to elements at indices 1, 2, 3
console.log("Sum [0, 2] after range update:", segTree.querySum(0, 2));  // Should be 13 (1 + 5 + 7)

// Test 3: Point update
segTree.pointUpdate(0, 10);  // Set element at index 0 to 10
console.log("Sum [0, 2] after point update:", segTree.querySum(0, 2));  // Should be 22 (10 + 5 + 7)

// Test 4: Combined operations
segTree.rangeUpdate(0, 5, 1);  // Add 1 to all elements
console.log("Value at index 2:", segTree.getValue(2));  // Should be 8 (7 + 1)
console.log("Sum of entire array:", segTree.querySum(0, 5));  // Sum of all elements after updates

// Test 5: Multiple range updates
segTree.rangeUpdate(2, 4, 3);  // Add 3 to elements at indices 2, 3, 4
segTree.rangeUpdate(1, 3, 2);  // Add 2 to elements at indices 1, 2, 3
console.log("Final sum [1, 3]:", segTree.querySum(1, 3));  // Sum after multiple updates

Segment Tree on 2D Matrix

Given a 2D matrix, we want to perform two types of operations efficiently:

1. Update: Change the value of a specific element in the matrix.
2. Range Query: Find the sum of elements in a submatrix defined by its top-left and bottom-right corners.

Segment Tree Structure

The 2D Segment Tree is a grid of Segment Trees, where each row has its own 1D Segment Tree to handle column operations. The main tree structure handles row operations.

Implementation

Let's see how to implement a Segment Tree for a 2D grid in JavaScript.

class SegmentTree2D {
  constructor(matrix) {
    this.rows = matrix.length;
    this.cols = matrix[0].length;
    this.tree = Array.from({ length: 2 * this.rows }, () => Array(2 * this.cols).fill(0));
    this.build(matrix);
  }

  build = (matrix) => {
    // Initialize leaves with the matrix elements
    for (let i = 0; i < this.rows; i++) {
      for (let j = 0; j < this.cols; j++) {
        this.tree[this.rows + i][this.cols + j] = matrix[i][j];
      }
    }

    // Build tree for each row
    for (let i = 0; i < this.rows; i++) {
      for (let j = this.cols - 1; j > 0; j--) {
        this.tree[this.rows + i][j] = this.tree[this.rows + i][2 * j] + this.tree[this.rows + i][2 * j + 1];
      }
    }

    // Build tree for the columns
    for (let i = this.rows - 1; i > 0; i--) {
      for (let j = 0; j < 2 * this.cols; j++) {
        this.tree[i][j] = this.tree[2 * i][j] + this.tree[2 * i + 1][j];
      }
    }
  }

  update = (row, col, value) => {
    let r = row + this.rows;
    let c = col + this.cols;
    this.tree[r][c] = value;

    // Update the row segment tree
    while (c > 1) {
      c = Math.floor(c / 2);
      this.tree[r][c] = this.tree[r][2 * c] + this.tree[r][2 * c + 1];
    }

    // Update the column segment tree
    while (r > 1) {
      r = Math.floor(r / 2);
      c = col + this.cols;
      while (c > 0) {
        this.tree[r][c] = this.tree[2 * r][c] + this.tree[2 * r + 1][c];
        c = Math.floor(c / 2);
      }
    }
  }

  query(row1, col1, row2, col2) {
    let sum = 0;

    // Adjust row range to include all leaves
    row1 += this.rows;
    row2 += this.rows + 1;

    while (row1 < row2) {
      if (row1 % 2 === 1) {
        sum += this.#queryRow(row1++, col1, col2);
      }
      if (row2 % 2 === 1) {
        sum += this.#queryRow(--row2, col1, col2);
      }
      row1 = Math.floor(row1 / 2);
      row2 = Math.floor(row2 / 2);
    }

    return sum;
  }

  #queryRow = (row, col1, col2) => {
    let sum = 0;
    col1 += this.cols;
    col2 += this.cols + 1;

    while (col1 < col2) {
      if (col1 % 2 === 1) {
        sum += this.tree[row][col1++];
      }
      if (col2 % 2 === 1) {
        sum += this.tree[row][--col2];
      }
      col1 = Math.floor(col1 / 2);
      col2 = Math.floor(col2 / 2);
    }

    return sum;
  }
}

// Example usage:
const matrix = [
  [1, 2, 3],
  [4, 5, 6],
  [7, 8, 9]
];

const segTree2D = new SegmentTree2D(matrix);
console.log(segTree2D.query(1, 1, 2, 2)); // Output: 28 (sum of submatrix [[5, 6], [8, 9]])
segTree2D.update(2, 2, 10);
console.log(segTree2D.query(1, 1, 2, 2)); // Output: 29 (sum of submatrix [[5, 6], [8, 10]])

Range Queries Problems

1D Range Queries

• Range Sum Query - Mutable
• Range Sum Query - Immutable

2D Range Queries

• Range Sum Query 2D - Mutable

Range Module

• Range Module

Counting Problems

• Count of Smaller Numbers After Self