Let’s Learn Segment Trees

It’s time I learned Segment Trees and taught : )

In the computer programming realm, the Segment trees are broadly used in cases where you have an array of elements and you are asked some queries on intervals.

Let’s say the number of elements N ~ 10⁵ , number of queries Q ~ 10⁵ when you take O(n) time to solve each query , YOU WILL NOT GET AN AC !!!. So, Let’s get that AC!!.

With Segment Trees we hope to bring down the complexity to O(log(n)), which is great. Let’s see the structure

segment tree for an 8 elements

In the tree, the root has an info on elements in interval [0,n) and each node, has 0 or two children. Left and right. If a node’s interval is [l, r) and l + 1 ≠ r, the interval of its children will be [l, mid) ( contains info of elements in this interval) and [mid, r) (contains info of elements in this interval ) in order where mid = (l+r)/2, so the height of this tree will be O(log(n)).

It would be easy to deal with segment tree, when it is 1-based. So, say the node is of index x , then if it has children the indexes in which they will be is 2x and 2x + 1. The main use cases of segment tree are as follows, for example info may be simply sum of elements, gcd of elements ,min , max …… You can get the idea of the info I am talking about. Let’s say we have an interval [x,y] and we need to find the nodes in the segment tree whose union of info is the info of elements in [x,y]. Below is the code to find these nodes.

vector<int> s;
void split(int x,int y, int id = 1,int l = 0, int r = n){// id is the index of the node
 if(x >= r or l >= y) return ; // in this case, intersect of [l,r) and [x,y) is empty
 if(x <= l && r <= y){
  s.push_back(id); 
  return ;
 }
 int mid = (l+r)/2;
 split(x,y,id * 2,l,mid);
 split(x,y,id * 2 + 1,mid,r);
}

Consider a simple problem, An array of elements a1,…..an and q queries. q types of queries.

1. Print sum of elements a_L to a_R
2. Update an element a_p to k

For the above problem, You may think you can solve it by using a cumulative array but you can’t because of the second query.

So let’s build the segment tree first. Observe the info here is sum of the elements. call build()

void build(int id = 1,int l = 0,int r = n){
 if(r - l < 2){ // l + 1 == r
  s[id] = a[l];
  return ;
 }
 int mid = (l+r)/2;
 build(id * 2, l, mid);
 build(id * 2 + 1, mid, r);
 s[id] = s[id * 2] + s[id * 2 + 1];
}

Now just call the above build, and the segment is ready. For the second query, You need to modify the value. Simple similar to the build above , you modify every node that has info about the element a[p] which is being updated to x. Just call modify(p,x)

void modify(int p,int x,int id = 1,int l = 0,int r = n){
	s[id] += x - a[p]; // modified info of the node.
	if(r - l < 2){	//	l = r - 1 = p
		a[p] = x;
		return ;
	}
	int mid = (l + r)/2;
	if(p < mid)
		modify(p, x, id * 2, l, mid);
	else
		modify(p, x, id * 2 + 1, mid, r);
}

For the first query Since we need the info of the elements in a range, You can use split as a backbone for different info. For the sum info call sum( l, r ), see code below.

int sum(int x,int y,int id = 1,int l = 0,int r = n){
	if(x >= r or l >= y)	return 0;
	if(x <= l && r <= y)	return s[id];
	int mid = (l+r)/2;
	return sum(x, y, id * 2, l, mid) +
	       sum(x, y, id * 2 + 1, mid, r);
}