I'm learning segment trees and their implementation. For the purpose, I tried solving the problem on CodeChef.
My tree implementation is as follows:
#include <bits/stdc++.h>
#define _MAX (1<<17)
#define MAX (_MAX<<1)
using namespace std;
struct node
{
int tails;
int heads;
void Merge(node &A,node &B) // Merge two nodes
{
tails=A.tails+B.tails;
heads=A.heads+B.heads;
}
void CreateLeaf() // Creates a leaf
{
tails=1;
heads=0;
}
};
node ST[MAX]; // Segment tree ST
/* res is the node storing the result
Information for range [l,r] is present in ST[idx]
Query is required for range [u,v] */
void Query(node &res,int l,int r,int u,int v,int idx)
{
if(l==u && r==v)
{
res=ST[idx];
return;
}
else
{
int mid=(l+r)/2;
if(mid>=v)
Query(res,l,mid,u,v,2*idx);
else if(mid<u)
Query(res,mid+1,r,u,v,2*idx+1);
else{
node left,right;
Query(left,l,mid,u,mid,2*idx); // left child is 2*idx
Query(right,mid+1,r,mid+1,v,2*idx+1); // right child is 2*idx+1
res.Merge(left,right);
}
}
}
void init(int idx,int l,int r) // initialises the tree
{
if(l==r)
{
ST[idx].CreateLeaf();
return;
}
else
{
int mid=(l+r)/2;
init(2*idx,l,mid);
init(2*idx+1,mid+1,r);
ST[idx].Merge(ST[2*idx],ST[2*idx+1]);
}
}
/* Updates the value of a single node.
When we flip the coins, no. of heads and tails get swapped */
void updatevalue(node &A)
{ swap(A.tails,A.heads); }
/* Updates the node ST[idx] and all the nodes descending the node ST[idx]
ST[idx] stores information for range [l,r] */
void update(int l,int r,int idx)
{
if(l==r) // returns as soon as we reach a leaf node
{
updatevalue(ST[idx]);
return;
}
else
{
updatevalue(ST[idx]);
int mid=(l+r)/2;
update(l,mid,2*idx); // Updates left child
update(mid+1,r,2*idx+1); // Updates right child
}
}
/* Range to be updated is [u,v]
ST[idx] stores information for range [l,r] */
void UpdateRange(int l,int r,int u,int v,int idx)
{
if(u<=l && r<=v)
{
update(l,r,idx);
// As soon as required node is found, update the node and all its descendant
return;
}
else
{
int mid=(l+r)/2;
if(mid>=v)
UpdateRange(l,mid,u,v,2*idx);
else if(mid<u)
UpdateRange(mid+1,r,u,v,2*idx+1);
else
{
UpdateRange(l,mid,u,mid,2*idx);
UpdateRange(mid+1,r,mid+1,v,2*idx+1);
}
ST[idx].Merge(ST[2*idx],ST[2*idx+1]);
}
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int n,q,x,y,z;
node res;
cin>>n>>q;
init(1,0,n-1); /* ST[1] stores information for [0,n-1] */
while(q--)
{
cin>>x>>y>>z;
if(x==1)
{
Query(res,0,n-1,y,z,1);
cout<<res.heads<<"\n";
}
else
{
UpdateRange(0,n-1,y,z,1);
}
}
return 0;
}
I've read many online sources and applied the segmented tree to the best of my knowledge, however my implementation is shown slow in accordance with time constraints for the problem. Please tell where my code is lacking. I feel my range update is a slow function but I can't find any method to speed it up as when we flip a coin, we have to update all ranges having that coin. Any better implementation for Range update or query is welcomed.