(Basic) Segment Trees

Problem Introduction§

Given an array of length $N$, we want to:

  • Query the maximum value of a range $[l, r]$
  • Change a particular index's value

For example, we want to be able to:

  • Query the minimum value in the array in the range $[1, 3]$
  • Then change the value of the second element (arr[2]) to 5

... and repeat that many times.

Possible approaches:

  • We could maintain the array as is - query time will be $O(N)$, but we can update in $O(1)$ time.
  • We could break the array into blocks of arbitrary length, say $k$, and maintain maximums for each of them; we will get $O(\frac{N}{k})$ query time (as we query up to $N/k$ blocks and $N \mod k$ individual elements) and $O(1)$ update time (as we only need to update the particular element and the whole block's maximum). This is the basic idea for a class of algorithms called square-root algorithms.
  • Use memoization to speed up queries. (We could even apply a bit of reasoning to drastically reduce the amount of queries we have to calculate; here's a hint: $[i, j] = \max{\left([i, i+2^n], [j - 2^n + 1, j]\right)}$) 1

Most data structures made for efficient max range queries need to be rebuilt every time the array changes, while breaking the array into blocks of arbitrary length generally gets us about $O(\sqrt{N})$ update time, which is good, but we can do better.


Oh, and maybe I should use some of Big O's siblings as well to mess with people.

Segment Tree with 2n elements§

We will assume that the number of elements in the array is a power of 2 for now, as it makes it easier to explain. Don't worry, we'll talk about the general case in another article.

Some of what we talked about as a possible approach (i.e. memoization and splitting queries) is a useful idea. We can decompose a query into smaller parts, calculate and store their answer, then when responding to a query, we can piece the answer together by stitching together ranges. When we have to update an element, we only need to update the answer for ranges that include the element, and none of the others. This is the main idea which underlies most query data structures.

The difference between Segment Trees, Square-Root Algorithms, and Sparse Tables is the way they split up ranges, which influences their time and memory complexity.

Concept§

Let us define an example array now to make things easier.

4 1 3 2 2 3 8 4 5 5 1 6 2 7 1 8
Example array

We start with the case of needing to find the maximum of a single index, e.g. $[2, 2]$ in which case, we already know that it's going to be arr[2].

Next, $[1, 2]$. We observe:

$$[1, 2] = \max{([1, 1], [2, 2])}$$

4 3 2 8 5 1 2 1 4 8 5 2 1 2 3 4 5 6 7 8
Partial Segment Tree

The division in ranges that defines segment trees is dividing ranges in half. It's a quite elegant idea really - we group two indices, then group two of those, and so on...

What about $[1, 6]$?

4 3 2 8 5 1 2 1 4 8 5 2 8 5 1 2 3 4 5 6 7 8
Partial Segment Tree

We can continue and create a binary tree, where each node has two children.

The nice thing about setting up a tree where nodes have a constant number of children (two) is that we can index children of node $i$ as $2 * i$ and $2*i + 1$. Conversely, the parent of $i$ would be $i/2$.

(Take a moment to convice yourself this works - this exploits the division implementation common in most languages - we always truncate the result, so (2*i)/2 and (2*i + 1)/2 are both equal to i)

From a different angle, we start building our tree with grouping every two elements. We define our segment tree as an array of size 16, such that the index of the parent of a child is i/2, and the index of the children is 2*i and 2*i + 1. Note that both (2*i)/2 and (2*i + 1)/2 are equal to i.

4 3 2 8 5 1 2 1 4 8 5 2 8 9 10 11 12 13 14 15 8 5 4 5 6 7 8 2 3 1
Segment Tree with indices

The key benefit to cutting ranges into smaller ranges this way is that we only need to edit one half of the range in case of an update: for example, to change the value of index $3$ to 9, we only need to modify $[3]$ and its ancestors ($[3, 4], [1, 4], [1, 8]$) and not bother with the other nodes.

This is the observation that leads to our $O(\log N)$ query and update time.

4 3 9 8 5 1 2 1 4 9 5 2 8 9 10 11 12 13 14 15 9 5 4 5 6 7 9 2 3 1
Segment Tree Change Propagation

Sidenote: Proof of Memory Complexity§

Our segment tree consists of levels which are half the size of the previous one (since there is a single parent node for every two child nodes). Let the number of base nodes be $N = 2^M$. Therefore, the total number of nodes is: $$ N + \frac{N}{2} + \frac{N}{4} + ... + 1 $$ or alternatively, $$ 2^{M} + 2^{M - 1} + 2^{M - 2} + ... + 2^{0} $$

We know, $$ 1 = \frac{1}{2} + \frac{1}{4} + \frac{1}{8} + ... + \frac{1}{\infty} $$ $$ \implies 2^0 = 2^{-1} + 2^{-2} + 2^{-3} + ... + 2^{-\infty} $$ $$ \implies 2^{M+1} = 2^{M} + 2^{M-1} + ... + 2^{0} + \underbrace{2^{-1} + 2^{-2} + 2^{-\infty}}_{\text{= 1}} $$ $$ \implies 2^{M+1} = 2^{M} + 2^{M-1} + ... + 2^{0} + 1 $$ $$ \implies 2^{M+1} > 2^{M} + 2^{M-1} + ... + 2^{0} $$

Therefore, our Segment Tree will consume at most $2^{M+1}$ memory, or in other words, $2 * N$ memory. ($2^{M+1}-1$ or $2 * N - 1$ to be precise)

See Zeno's Paradox for more detail; specifically that the $1 = \frac{1}{2} + \frac{1}{4} + \frac{1}{8} + ...$ does indeed converge.

Sidenote: Segment Tree Power!!!§

Segment Trees are an extremely flexible data structure - we could also build a segment tree for sum queries. In fact, we can build a segment tree for any operation which is associative (i.e. order of operations doesn't matter; $(a \centerdot b) \centerdot c = a \centerdot (b \centerdot c)$).

This property ensures that we can prevaluate range answers and then combine them for the right answer. We can be certain that $$ \left([1] \centerdot [2]\right) \centerdot \left([3] \centerdot [4]\right) \centerdot [5] = [1] \centerdot [2] \centerdot [3] \centerdot [4] \centerdot [5] $$

Now that we're finished talking about the basic concept of the segment tree, on to implementing it!


We define the following functions:

  • $\text{build} (\text{arr})$ - actually constructs our segment tree from the array.
  • $\text{query} (l, r)$ - returns the maximum element in $[l, r]$ range of the array.
  • $\text{update} (\text{value, index})$ - set element at $\text{index}$ of the array to $\text{value}$. (i.e. set arr[index] = value)

Implementation Note: We'll be using the max operation here. The code present below assume commutativity, though supporting non-commutative operations is as simple as ensuring order of operands when modifying the tree. (See my generic version which does not assume commutativity.)

Building build§

We observe that the indices for the ranges of $[1,1]...[i, i]...[N, N]$ (i.e. the base cases) in the segment tree are $[N, N]...[N + i, N + i]...[2 * N - 1, 2 * N - 1]$, where $N$ is the total number of elements. (Refer to the previous section for an understanding of as to why the total number of nodes in a binary tree are $2*N$)

In case you've forgotten, the number of elements in a both sides inclusive range $[l, r]$ is $r - l + 1$.

We fill the second half of the array with our initial elements, then start working our way down. Make sure that you understand why buf[i] is calculated from buf[2*i] and buf[2*i + 1].

code/segtree-two-power/src/max_segtree.rs
10pub struct SegTree<const N: usize>
11where
12 [u32; 2*N]: Sized,
13{
14 buf: [u32; 2*N],
15}
16
17impl<const N: usize> SegTree<N>
18where
19 [u32; 2*N]: Sized,
20{
21 pub fn new (arr: &[u32; N]) -> Self {
22 let mut buf= [u32::default(); 2*N];
23 buf[N..2*N].copy_from_slice(arr);
24 for i in (1..N).rev() {
25 buf[i] = <u32>::max(buf[2*i], buf[2*i + 1]);
26 }
27
28 SegTree {buf}
29 }
30

Querying query§

For querying, we start with $[l, l]$ and $[r, r]$ at the lowest level, then move up the tree, ensuring at all times that both of the ranges are subsets of the original query range, $[l, r]$.

We start at $[l, l]$, then merge in rightward ranges moving up to eventually get to $[l, r]$. Similarly, we start at $[r, r]$ and then merge in leftward ranges moving up to eventually get to $[l, r]$.

One observation that we can make is that determining whether a range is a "left child" or a "right child" is that all "left children" have odd indices, while all "right children" have even ones.

A range, say, $[1, 2]$ will have two children: $[1, 1]$ and $[2, 2]$. Out of these, $[1, 1]$ will be the "left child" and $[2, 2]$ will be the right child. When we're at a range, say $[1, 1]$, we can use the index of said range to determine whether merging $[1, 1]$'s parent will extend our overall range leftward or rightward. In this case, the index of $[1, 1]$ will be odd, which means that merging in $[1, 2]$ will extend our range rightward.

For example:

The parent of $[l, l]$ may include $[l - 1, l]$ or $[l, l + 1]$, depending on whether $l$ is even or odd. We always want to move towards $[l, l + 1]$, since $l - 1$ is NOT in our range, but $l + 1$ is.

In the case that $[l, l]$'s parent is $[l - 1, l]$, we move up to $[l + 1, l + 2]$ and merge $[l, l]$ in.

Similarly, if $[r, r]$'s parent is $[r, r + 1]$, we move up to $[r - 2, r - 1]$ and merge $[r, r]$ in.

Keep in mind that $[l, r] \equiv \lbrace l, l + 1, l + 2, ..., r \rbrace$

4 3 2 8 5 1 2 1 4 8 5 2 8 9 10 11 12 13 14 15 8 5 4 5 6 7 8 2 3 1
Querying $[1, 7]$
code/segtree-two-power/src/max_segtree.rs
31 pub fn query (&self, (mut l, mut r): (usize, usize)) -> u32 {
32 l += N - 1; r += N - 1;
33 let mut ans = self.buf[l];
34 while l <= r {
35 if l % 2 == 1 {
36 ans = <u32>::max(ans, self.buf[l]); l += 1;
37 }
38 if r % 2 == 0 {
39 ans = <u32>::max(ans, self.buf[r]); r -= 1;
40 }
41 l /= 2; r /= 2;
42 }
43 ans
44 }

Updating update§

Updating is straightforward - we start at $[i, i]$, then divide the index by half to move up a level and update that range. Repeat until all ranges which include $i$ have been updated.

4 3 9 8 5 1 2 1 4 9 5 2 8 9 10 11 12 13 14 15 9 5 4 5 6 7 9 2 3 1
Setting $[3] = 9$
code/segtree-two-power/src/max_segtree.rs
46 pub fn update (&mut self, val: u32, mut idx: usize) {
47 idx += N - 1;
48 self.buf[idx] = val;
49 idx /= 2;
50
51 while idx != 0 {
52 self.buf[idx] = <u32>::max(self.buf[2*idx], self.buf[2*idx + 1]);
53 idx /= 2;
54 }
55 }

Segment Tree with a general amount of elements§

The ideas explored above can also be applied to arrays of any size with only slight changes. However, that is a topic for another day.

Further Reading/References/Footnotes§



As always, feel free to write to me to point out an error, suggest a topic, or just say hi!