Segment tree implementation

Question

I'm learning segment trees and their implementation. For the purpose, I tried solving the problem on CodeChef.

Full code here

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.

Answer 1

In your if-else logic :

• you can get rid of a few returns

• if you have something like if (condition) { some_code(); some_more_code1(); } else { some_code(); some_more_code2(); }, you can write it by moving the common code : some_code() if (condition) { some_more_code1(); } else { some_more_code2(); } (assuming that the code doesn't affect the evaluation of the condition).

• if you have a test with an empty then block, you can rewrite it with the negation to have a test with no else block.

Your CreateLeaf() seems to be something that could be done in a constructor.

Your updatevalue() seems to be something that could be done in the node class.

You can move the declaration of int x, y, z to a smaller scope.

Query should probably return a value instead of taking a reference.

Merge could take const references.

Using using namespace std is usually frowned upon.

Taking (most of) these comments (and some others) into account, I have :

#include <bits/stdc++.h>

#define _MAX (1<<17)
#define MAX (_MAX<<1)
using namespace std;

class Node
{
    public:
    Node(): tails(1), heads(0) {}

    void Merge(const Node &A,const Node &B) // Merge two Nodes
    {
        tails=A.tails+B.tails;
        heads=A.heads+B.heads;
    }
    /* Updates the value of a single Node.
   When we flip the coins, no. of heads and tails get swapped */
    void Swap()
    {
        swap(tails,heads);
    }
    int getNbHeads() { return heads; }
    private:
    int tails;
    int heads;
};


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]  */

Node Query(int l,int r,int u,int v,int idx)
{
    if(l==u && r==v)
    {
        return ST[idx];
    }
    else
    {
        int mid=(l+r)/2;
        if(mid>=v)
            return Query(l,    mid,u,v,2*idx);
        else if(mid<u)
            return Query(mid+1,r,  u,v,2*idx+1);
        else{
            Node res;
            res.Merge(
                Query(l,    mid,u,    mid,2*idx), // left child is 2*idx
                Query(mid+1,r,  mid+1,v,  2*idx+1) // right child is 2*idx+1
            );
            return res;
        }
    }
}

void init(int idx,int l,int r) // initialises the tree
{
    if(l!=r)
    {
        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 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)
{
    ST[idx].Swap();
    if(l!=r)
    {
        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)
    {
        // As soon as required Node is found, update the Node and all its descendant
        update(l,r,idx);
    }
    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;
    cin>>n>>q;
    init(1,0,n-1); /* ST[1] stores information for [0,n-1] */
    while(q--)
    {
        int x, y, z;
        cin>>x>>y>>z;
        if(x==1)
        {
            cout<< Query(0,n-1,y,z,1).getNbHeads()<<"\n";
        }
        else
        {
            UpdateRange(0,n-1,y,z,1);
        }
    }
    return 0;
}

Also, I forgot to tell it but it would probably be worth encapsulating the logic around ST in some kind of class or at least pass the array as an argument as global variable can make things hard to track/maintain.