Advanced Segment Trees

Written by Johannes Kapfhammer.

Contents

Recursive Segment Tree
Lazy Propagation
Examples
- Min Query, Set Updates
- Sum Query, Add Updates
Full Code
Iterative Lazy Segment Trees

This article assumes that you’ve already read the article about segment tree basics. Whereas the linked article uses an iterative implementation of a segment tree, we opt for a recursive implementation. Implementing queries and updates recursively is more natural, as we’ll see (it fits the style of divide and conquer, and there are no bit tricks involved), and it makes lazy updates simpler to implement.

Recursive Segment Tree

We first start with a regular segment tree, but implemented recursively. As we’ve seen in the basic article, we should think about segment trees as operating on abstract data instead of some particular hard-coded operation. Formally, a segment tree allows range queries $x_i\circ x_{i+1} \circ x_{i+2} \circ \dots \circ x_{j-1}$ for any range $[i, j)$ and an arbitrary (but fixed) associative operation $\circ$ . Associativity means that we must have $a\circ (b\circ c)=(a\circ b)\circ c$ for any $a$ , $b$ and $c$ .

In case of regular segment trees, the supported element updates are changing $x_i$ into $f(x_i)$ for arbitrary functions $f$ .

We model this as follows:

struct Value { ... }; // the value we store at each node
struct Update { ... }; // information for an update
Value combine(Value, Value) { ... }
Value identity_value() { ... }
Value apply_update(Update, Value) { ... }

Here, the structure Value is the type of the $x_i$ . The operation $\circ$ is the function combine.

Additionally, we require there to be an identity element $e$ , for which we have $e\circ x=x$ and $x\circ e=x$ for all $x$ . Note that it is very easy to support such an element (just set $e$ to a value that is not valid and check in $\circ$ whether one element is equal to $e$ and if yes return the other one). Having such an identity element will simplify the code.

Finally, function $f$ is stored as structure Update and the result of the function call $f(x)$ can be computed using the function apply_update. Note that storing updates in a separate structure is overkill in the case of regular segment trees (we could just say that $x_i$ is replaced by $x_i'$ ), but we’ll need that for generalizing it to lazy segment trees.

We store the values in a heap: The root is stored at position $1$ . The children of the node at position $pos$ are at positions $2\cdot pos$ and $2\cdot pos+1$ . The parent of $pos$ is $\lfloor pos/2\rfloor$ .

The two main functions are update_ and query_. Both take as parameters ql (query left) and qr (query right), the range $[ql, qr)$ of the query (from ql inclusively to qr-1 inclusively``). The parameter pos is the index of the node in the heap. This node stores the value of $q_l\circ q_{l+1}\circ \dots\circ q_{r-1}$ , and those $l$ and $r$ are given as parameter.

As the order of the parameters can be tricky when calling, we make them private and wrap them into functions update and query that only take $ql$ and $qr$ and find the right initial values for $pos$ and $l$ and $r$ .

The two helper functions apply and recompute are one-liners that apply an update to a current node or recompute the value from its two children. Those functions will get more complicated later when we do lazy updates.

// we want the size of the segment tree to be a power of two
int next_power_of_two(int x) {
  int p2 = 1;
  while (p2 < x)
    p2 *= 2;
  return p2;
}

struct SegmentTree {
  int n;
  vector<Value> tree;

  // build segtree of size at least min_n
  SegmentTree(int min_n)
    : n(next_power_of_two(min_n)),
      tree(2*n, identity_value()) {}

  // combines all values in range [l, r)
  Value query(int l, int r) {
    assert(0 <= l && l < r && r <= n);
    return query_(l, r, 1, 0, n);
  }

  // updates all values in range [i, i+1)
  // once we implement lazy propagation, this will allow arbitrary ranges [l, r)
  void update(int l, int r, Update a) {
    assert(r - l == 1); // only allow element updates (for now)
    assert(0 <= l && l < r && r <= n);
    return update_(a, l, r, 1, 0, n);
  }

private:
  // applies the update to the current node
  void apply(int pos, Update a) {
    tree[pos] = apply_update(a, tree[pos]);
  }

  // recomputes the value of position "pos"
  void recompute(int pos) {
    tree[pos] = combine(tree[2*pos], tree[2*pos+1]);
  }

  Value query_(int ql, int qr, int pos, int l, int r) {
    // completely contained: return value
    if (ql <= l && r <= qr) { return tree[pos]; }
    // not overlapping: return nothing
    if (qr <= l || r <= ql) { return identity_value(); }
    // otherwise: recurse
    int m = (l+r)/2;
    Value ans_l = query_(ql, qr, 2*pos, l, m);
    Value ans_r = query_(ql, qr, 2*pos+1, m, r);
    return combine(ans_l, ans_r);
  }

  void update_(Update a, int ql, int qr, int pos, int l, int r) {
    // at leaf node: apply update
    if (ql <= l && r <= qr) { apply(pos, a); return; }
    // not overlapping: do nothing
    if (qr <= l || r <= ql) { return; }
    // otherwise: recurse
    int m = (l+r)/2;
    update_(a, ql, qr, 2*pos, l, m);
    update_(a, ql, qr, 2*pos+1, m, r);
    recompute(pos); // update the value of the current position
  }
};

In case you aren’t familiar with C++, the constructor is initializing the class members using an initializer list:

// build segtree of size at least min_n
SegmentTree(int min_n)
  : n(next_power_of_two(min_n)),
    tree(2*n, identity_value()) {}

This code is equivalent to

// build segtree of size at least min_n
SegmentTree(int min_n) {
  n = next_power_of_two(min_n);
  tree = vector<Value>(2*n, identity_value());
}

It is recommended to use initializer lists whenever possible, so the article will use them throughout.

Lazy Propagation

Now we want to support arbitrary range updates. The reason we can’t support them with our segment tree from before is because of the first line of the update_ function: In case we are not at a leaf node, we don’t know how to update the whole range.

// at leaf node: apply update
if (ql <= l && r <= qr) { apply(pos, a); return; }

So how can we make apply work for arbitrary nodes? The trick is as follows: We just “pretend” we do it, but actually we just place a note in the node that we will do it later. Only when asked for the value of the node or one of its children, we really execute the operation.

More formally, say our update is function $f$ and we want to change $x_i$ into $f(x_i)$ for all $i\in [l,r)$ for arbitrary ranges $[l,r)$ . In general, we would have to apply $f$ for every element. However, if we have the property that $f(x\circ y)=f(x)\circ f(y)$ (distributivity), we can compute the result of $f$ on a range without first having to apply it to each element. The trick from before was to only compute the result if necessary.

How to implement it? We need another function in our segment tree called propagate. Propagate applies the pending updates of a node to its children and removes the note that those updates need to be executed.

Why is this fast? Consider an arbitrary query. We build our answer out of $\mathcal O(\log n)$ nodes that together cover our query range, because we only walk down two paths from root to some leaf. So we only need to execute $\mathcal O(\log n)$ propagates.

But why is propagate fast? When we do update $f$ after update $g$ on a range (i.e. changing $x$ to $f(g(x))$ ) and we would store both $f$ and $g$ in our node, our representation will explode and we have slow updates again.

We resolve this by requiring another operation $\otimes$ that allows us to efficiently chain updates. This operation needs to fulfill the properties that $(f\otimes g)(x)=f(g(x))$ (it’s the same as chaining), and $(f\otimes g)\otimes h=f\otimes (g\otimes h)$ (associativity).

Also, it will be convenient for us to have a function $\operatorname{id}$ which has the property that $f\otimes \operatorname{id}=f$ , $\operatorname{id}\otimes f=f$ for all updates $f$ and $\operatorname{id}(x)=x$ for all values $x$ .

Alltogether, our interface for the segment tree will look like this:

struct Value { ... };
Value identity_value() { ... }
Value combine(Value a, Value b) { ... }

struct Update { ... };
Update identity_update() { ... }
Value apply_update(Update a, Value x) { ... }
Update combine_updates(Update a, Update b) { ... }

The only changes required to our segment tree are an additional vector lazy to store the updates, a function propagate that pushes down the updates, and combining two updates in the apply function. Also we need to call propagate at the right place in query_ and update_:

It is important to understand what the value of lazy[pos] conceptually means: It is the operation that is yet to be performed on the subtree of pos, but has already been applied to tree[pos].

// store the updates in a separate vector (initialized to have size 2*n)
vector<Update> lazy;

// pushes lazy values down the tree
void propagate(int pos) {
  apply(2*pos, lazy[pos]);
  apply(2*pos+1, lazy[pos]);
  lazy[pos] = identity_update();
}

// applies the update to the current node
void apply(int pos, Update a) {
  tree[pos] = apply_update(a, tree[pos]); // same as before
  lazy[pos] = combine_updates(a, lazy[pos]); // new: mark operation for subtree
}

Value query_(int ql, int qr, int pos, int l, int r) {
  ... // (checks same as before)
  // otherwise: recurse
  propagate(pos);
  ... // (rest is same as before)
}

void update_(Update a, int ql, int qr, int pos, int l, int r) {
  ... // (checks same as before)
  // otherwise: recurse
  propagate(pos);
  ... // (rest is same as before)
}

Examples

Let’s see how we can model typical updates with operations $\circ$ and $\otimes$ .

Min Query, Set Updates

Query: compute the minimum of a range

Update: set a range to a constant value

Our elements are integers. We define $x\circ y$ as $\min(x,y)$ because we’re interested in the minimum. Updates to set to a value $c$ are defined as $f_c(x)=c$ .

Let’s check whether these operations are good:

$\min(\min(x,y),z)=\min(x,\min(y,z))$ : the order in which we take the minimum doesn’t matter
$\min(x,\infty)=\min(\infty,x)=x$ : We need an element $\infty$ , which can never be smaller than a real value.
$f_u(\min(x,y))=\min(f_u(x),f_u(y))$ : This is the crucial property we need. When we set a range to a value, the minimum will also change to that.
$f_u\otimes f_v=f_u$ : setting to $v$ and then setting to $u$ is the same as directly setting to $u$ .
$f_u(f_v(x))=f_u(v)=u$ and $(f_u\otimes f_v)(x)=f_u(x)=u$ , therefore $f_u(f_v(x))= (f_u\otimes f_v)(x)$ .
$f_u\otimes (f_v\otimes f_w)=f_u=f_u\otimes f_w=(f_u\otimes f_v)\otimes f_w$ (associativity)
We define $f_{-1}(x)=x$ because setting to $-1$ is not needed and make sure this acts as the neutral element.

Thus our code would look like this:

struct Value { int x; };
Value identity_value() { return { (int)1e9+1 }; } // 1e9+1 is larger than any real value
Value combine(Value a, Value b) { return { min(a.x, b.x) }; }

struct Update { int v; };
Update identity_update() { return { -1 }; } // we never set to -1
Value apply_update(Update a, Value x) {
  if (a.v == -1) return x; // check for identity
  return { a.v; };
}
Update combine_updates(Update a, Update b) {
  if (a.v == -1) return b; // check for identity
  return a;
}

Sum Query, Add Updates

Query: compute the sum of a range

Update: add a constant value to all elements of a range

If we define this as before just with $+$ instead of $\min$ , and $f_v(x)=x+v$ , it will not work because of the distributivity: $f_u(x+y)=x+y+u$ , but $f_u(x)+f_u(y)=x+y+2u$ .

The key insight here is that we need the length of the range for the update. We define our elements as pair $(\text{value}, \text{length})$ , where for our leaf values $\text{length}=1$ . Then:

$(x,k)\circ(y,l)=(x+y,k+l)$ (obviously associative, with neutral element $(0,0)$ )
$f_v((x,k))=(x+v\cdot k, k)$ : Since we add $v$ to each element in the range, the sum gets increased by $v\cdot k$ .
$f_v((x,k)\circ(y,l))=((x+y)+v\cdot (k+l),k+l)=(x+v\cdot k,k)\circ(y+v\cdot l,l)=f_v((x,k))\circ f_v((y,l))$ : Now, we have distributivity!
$f_u\otimes f_v=f_{u+v}$ : adding $u$ and then adding $v$ is the same as adding $u+v$ (and addition is obviously associative).
$f_{0}((x,k))=(x,k)$ : we don’t need a special value for our neutral element.

struct Value { int x, k; };
Value identity_value() { return { 0, 0 }; }
Value combine(Value a, Value b) { return { a.x+b.x, a.k+b.k }; }

struct Update { int v; };
Update identity_update() { return { 0 }; }
Value apply_update(Update a, Value x) { return {x.x + a.v*x.k}; }
Update combine_updates(Update a, Update b) { return { a.v + b.v }; }

Full Code

For reference, the full code of the segment tree with lazy propagation is shown below:

// shorter implementation using std::__lg which computes the logarithm in base 2
int next_power_of_two(unsigned x) { return 1<<__lg(x-1)+1); }

struct SegmentTree {
  int n;
  vector<Value> tree;
  vector<Update> lazy;

  // build segtree of size at least min_n
  SegmentTree(int min_n)
    : n(next_power_of_two(min_n)),
      tree(2*n, identity_value()),
      lazy(2*n, identity_update()) {}

  // build segtree on an array of initial values
  SegmentTree(vector<Value> const& base)
    : n(next_power_of_two(base.size())),
      tree(2*n, identity_value()),
      lazy(2*n, identity_update()) {
    for (int i=0; i<(int)base.size(); ++i)
      tree[n+i] = base[i];
    build(1, 0, n);
  }

  // combines all values in range [l, r)
  Value query(int l, int r) {
    assert(0 <= l && l < r && r <= n);
    return query_(l, r, 1, 0, n);
  }

  // updates all values in range [l, r)
  void update(int l, int r, Update a) {
    assert(0 <= l && l < r && r <= n);
    return update_(a, l, r, 1, 0, n);
  }

private:
  // applies the update to the current node
  void apply(int pos, Update a) {
    tree[pos] = apply_update(a, tree[pos]);
    lazy[pos] = combine_updates(a, lazy[pos]);
  }

  // recomputes the value of position "pos", assuming lazy[pos]==identity_update()
  void recompute(int pos) {
    tree[pos] = combine(tree[2*pos], tree[2*pos+1]);
  }

  // pushes lazy values down the tree
  void propagate(int pos) {
    apply(2*pos, lazy[pos]);
    apply(2*pos+1, lazy[pos]);
    lazy[pos] = identity_update();
  }

  // build segtree assuming only leaf nodes are correct
  void build(int pos, int l, int r){
    if (r - l == 1) // leaf: do nothing
      return;
    int m = (l+r)/2;
    build(2*pos, l, m);
    build(2*pos+1, m, r);
    recompute(pos);
  }

  Value query_(int ql, int qr, int pos, int l, int r) {
    // completely contained: return value
    if (ql <= l && r <= qr) { return tree[pos]; }
    // not overlapping: return nothing
    if (qr <= l || r <= ql) { return identity_value(); }
    // otherwise: recurse
    propagate(pos);
    int m = (l+r)/2;
    Value ans_l = query_(ql, qr, 2*pos, l, m);
    Value ans_r = query_(ql, qr, 2*pos+1, m, r);
    return combine(ans_l, ans_r);
  }

  void update_(Update a, int ql, int qr, int pos, int l, int r) {
    // completely contained: update lazy
    if (ql <= l && r <= qr) { apply(pos, a); return; }
    // not overlapping: do nothing
    if (qr <= l || r <= ql) { return; }
    // otherwise: recurse
    propagate(pos);
    int m = (l+r)/2;
    update_(a, ql, qr, 2*pos, l, m);
    update_(a, ql, qr, 2*pos+1, m, r);
    recompute(pos);
  }
};

Iterative Lazy Segment Trees

Note that it is possible to do lazy segment trees iteratively. As you don’t need the recursion stack, the code will be more efficient, which can be useful for e.g. online competitions where you can copy/paste already written code. During SOI and IOI-style competitions where you don’t have any pre-written library available, I wouldn’t recommend it though as it can be tricky to get right.

In any case, you may want to read up on the details here: https://codeforces.com/blog/entry/18051