Segment Trees

Segment Trees is a heap-like data structure used for making range update / query operations in logarithmic time. It stores information about array intervals in a tree. The Segment Tree can be easily generalized to larger dimensions. For instance, with a two-dimensional Segment Tree you can answer sum or minimum queries over some subrectangle of a given matrix in only $O(\log^2 n)$ time.

Queries can be:

Range Sum Queries: Find the sum of elements in a given range.
Range Minimum/Maximum Queries: Find the minimum or maximum element in a given range.
Range GCD/LCM Queries: Compute the GCD or LCM of elements in a range.

Structure

Segment tree is a binary Tree where:

Each leaf node represents a single element of the array.
The internal node represents a function over a segment of the array.

The root of the tree is the segment $a[0 \dots n-1]$. We then split the array into 2 halves left child will have $a[0 \dots n/2 -1]$ and right child will have $a[n/2 \dots n-1]$. This is repeated until the segment reaches size 1.

graph TD;
    A[0...6];
    B[0...2]
    C[3...6]
    D[0...0]
    E[1...2]
    F[3...4]
    G[5...6]
    H[1...1]
    I[2...2]
    J[3...3]
    K[4...4]
    L[5...5]
    M[6...6]

    A --> B;
    A --> C;
    B --> D;
    B --> E;
    C --> F;
    C --> G;
    E --> H;
    E --> I;
    F --> J
    F --> K
    G --> L
    G --> M

Build

We start at the bottom level, leaf nodes and assign them their respective values. We can compute the previous level from current level, using merge function. Repeat this until we reach the root node.

Recursively, we start at the root node, and each node constructs 2 child nodes and merges the computed values in these children.

We use an array to store the tree. The left child of a node at index $i$ is at $2i$ and the right child at $2i + 1$.

val tree = MutableList<Int>(4*n)
fun build(arr: List<Int>, tree: MutableList<Int>, idx: Int, l: Int, r: Int){
    if (l == r){
        tree[idx] = arr[l]
    } else {
        val m = (l + r) / 2
        build(arr, tree, idx*2, l, m) // left child
        build(arr, tree, idx*2 + 1, m + 1, l) // right child
        tree[idx] = tree[2*idx] + tree[2*idx + 1]
    }
}

Efficient implementation

Leaf nodes start at index $n$ and will go upto $2*n-1$. To get parents we will start at index n-1 and move upwards.

fun build(arr: List<Int>): List {
    val n = arr.size
    val tree = MutableList <Int>(2*n)

    arr.forEachIndexed{ idx, it ->
        tree[n + idx] = it
    }

    for(i in n-1 downTo 0) {
        tree[i] = tree[i shl 1] + tree[i shl 1 + 1]
    }
}

Queries

Sum queries

Input is 2 integers $l$ and $r$. We have to compute the sum of $a[l \dots r]$ in O(log n). We traverse the segment tree and use the precomputed sum. While traversing the tree if we are at a node that covers $a[x \dots y]$ there are 3 cases:

Equals: $x \dots y ~=~ l \dots r$
In left or right child: $ x \le l \le r \le m$ or $m \le l \le r \le y$ where m = $\frac{x + y}{2}$
In the intersection of left & right child.

fun sum(tree:List<Int>, l: Int, r:Int): Int {
    var res = 0
    val n = tree.size / 2
    var li = l + n
    var ri = r + n

    while(l < r) {
        if(l and 1 == 1) res += tree[li++]
        if(r and 1 == 1) res += tree[--ri]
        li = li shr 1
        ri = ri shr 1
    }
}

Updates

Traverse until the leaf node and update all nodes in the path.

fun update(tree: List<Int>, idx: Int, newVal: Int) {
    val n = tree.size / 2
    tree[n + idx] = newVal

    var i = n + idx
    while(i > 1) {
        tree[i shr 1] = tree[i] + tree[i xor 1]
        i = i shr 1
    }
}