Implementation of Fenwick Tree gives TLE

81 Views Asked by At

Link to Code Pastebin

#include <bits/stdc++.h>

using namespace std;

vector <long long int> update(int x,long long int val,vector <long long int> b,int n)
{
    b[x]+=(val);
    x+=(x&-x);
    while(x<=n)
    {
        b[x]+=(val);
        x+=(x&-x);
    }
    return b;
}

long long int getsum(int i,int j,vector <long long int> b)
{
    i=i-1;
    long long int sum1=b[i];
    i-=(i&-i);
    while(i>0)
    {
        sum1+=(b[i]);
        i-=(i&-i);
    }
    long long int sum2=b[j];
    j-=(j&-j);
    while(j>0)
    {
        sum2+=b[j];
        j-=(j&-j);
    }
    return(sum2-sum1);
}

int main()
{
    int n,q;
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cin>>n>>q;
    vector <int> h(n);
    for(int i=0;i<n;i++)
    {
        cin>>h[i];
    }
    vector <long long int> a(n);
    for(int i=0;i<n;i++)
    {
        cin>>a[i];
    }
    vector <int> Pl(n,-1);
    vector <int> Pr(n,-1);
    stack <int> s;
    s.push(0);
    for(int i = 1; i < n; i++)
    {
    
        if (s.empty())
        { 
            s.push(i);
            continue; 
        } 
        while (s.empty() == false && h[s.top()] < h[i]) 
        {          
            Pl[s.top()] = i;
            s.pop(); 
        }
        s.push(i); 
    }
    while(s.empty()==false)
    {
        s.pop();
    }
    
    s.push(n-1);
    // iterate for rest of the elements 
    for(int i = n-2; i >= 0; i--)
    {
        if (s.empty())
        { 
            s.push(i);
            continue; 
        }
        while (s.empty() == false && h[s.top()] < h[i]) 
        {          
            Pr[s.top()] = i;
            s.pop(); 
        } 
        s.push(i); 
    }
    while(s.empty()==false)
    {
        s.pop();
    }
    vector <pair <int,int>> timel(n);
    vector <int> arrayl;
    s.push(0);
    int t = 1;
    timel[0].first = t;
    t++;
    arrayl.push_back(0);
    arrayl.push_back(a[0]);
    int i = 1;
    while(i<n)
    {
        if(s.empty())
        {
            s.push(i);
            timel[i].first = t;
            t++;
            arrayl.push_back(a[i]);
            i++;
            if(i==n)
                break;
        }
        if(h[i]<h[s.top()])
        {
            s.push(i);
            timel[i].first = t;
            t++;
            arrayl.push_back(a[i]);
            i++;
        }
        else
        {
            timel[s.top()].second = t;
            t++;
            arrayl.push_back(-a[s.top()]);
            s.pop();
        }
    }
    while(s.empty()==false)
    {
        timel[s.top()].second = t;
        t++;
        arrayl.push_back(-a[s.top()]);
        s.pop();
    }
    vector <pair <int,int>> timer(n);
    vector <int> arrayr;
    s.push(n-1);
    t = 1;
    timer[n-1].first = t;
    t++;
    arrayr.push_back(0);
    arrayr.push_back(a[n-1]);
    i = n-2;
    while(i>=0)
    {
        if(s.empty())
        {
            s.push(i);
            timer[i].first = t;
            t++;
            arrayr.push_back(a[i]);
            i--;
            if(i==0)
                break;
        }
        if(h[i]<h[s.top()])
        {
            s.push(i);
            timer[i].first = t;
            t++;
            arrayr.push_back(a[i]);
            i--;
        }
        else
        {
            timer[s.top()].second = t;
            t++;
            arrayr.push_back(-a[s.top()]);
            s.pop();
        }
    }
    while(s.empty()==false)
    {
        timer[s.top()].second = t;
        t++;
        arrayr.push_back(-a[s.top()]);
        s.pop();
    }
    vector <long long int> bitr(arrayr.size(),0);
    vector <long long int> bitl(arrayl.size(),0);
    for(int i=1;i<arrayl.size();i++)
    {
        bitl = update(i,arrayl[i],bitl, arrayl.size()-1);
        bitr = update(i,arrayr[i],bitr, arrayr.size()-1);
    }
    for(int i=0;i<q;i++)
    {
        int type;
        cin>>type;
        if(type==2)
        {
            int s,d;
            cin>>s>>d;
            s--;
            d--;
            int x = s;
            int y = d;
            if(s>d)
            {
                while (Pl[x] != Pl[y])
                {
                    if (x < y)
                        x = Pl[x];
                    else
                        y = Pl[y];
                }
                if(x != s)
                {
                    cout<<-1<<endl;
                }
                else
                {
                    long long int ans = getsum(timer[s].first,timer[d].first,bitr);
                    cout<<ans<<endl;
                }
            }
            else
            {
                while (Pr[x] != Pr[y])
                {
                    if (x > y)
                        x = Pr[x];
                    else
                        y = Pr[y];
                }
                if(y != s)
                {
                    cout<<-1<<endl;
                }
                else
                {
                    long long int ans = getsum(timel[s].first,timel[d].first,bitl);
                    cout<<ans<<endl;
                }
            }
        }
        else
        {
            int s,k;
            cin>>s>>k;
            long long int val = k-a[s-1];
            bitl = update(timel[s-1].first,val,bitl,arrayl.size()-1);
            bitl = update(timel[s-1].second,-val,bitl,arrayl.size()-1);
            bitr = update(timer[s-1].first,val,bitr,arrayr.size()-1);
            bitr = update(timer[s-1].second,-val,bitr,arrayr.size()-1);
            a[s-1] = k;
        }
    }
    return 0;
}

Sample Input:
5 6
3 1 4 1 5
1 2 4 8 16
2 5 2
1 3 5
2 3 4
2 1 4
2 5 1
2 3 3


Sample Output:
22
13
-1
22
5

I am trying this question where I have implemented a Fenwick tree to calculate the sum of the values of nodes on a path of the tree. The Fenwick tree has been implemented on the DFS array which stores the node value as positive during "in" time and negative during "out" time [(Link to the referred algorithm)](https://codeforces.com/blog/entry/78564).

There is an edge from Node i to Node j if h[i] > h[j] (heights) and there can be no edges which change direction i.e. the edges have to be added in either increasing or decreasing i. The values of each node is stored in vector a and can be updated. So I have constructed 2 trees from both directions.

An image of sample input and sample edges

The program runs well for the sample test cases. But it gives TLE even on the constraints N, Q <= 1000. The original constraints of the problem is N, Q <= 2*105

Pl vector stores the parent of tree rooted from left to right. Pr stores parent of the tree rooted on right to left. Similarly I have created arrayl and arrayr (DFS arrays), timer and timel (start and end time array).

The queries are of two types:
Either ask if there exists a path, then find the maximum sum of nodes.
Or Update the value of a certain node.

I calculated the order of the program to be about O((N+Q)logN). Am I wrong? Where am I going wrong in the code?

References:
Hill climbing problem requiring Dynamic programming
https://www.geeksforgeeks.org/next-greater-element/
https://medium.com/@adityakumar_98609/fenwick-tree-binary-index-tree-aca7824d9c2a

Add comments in case of further clarifications. I have posted as many relevant information as possible.

0

There are 0 best solutions below