HackerEarth Tree and K-Tuple problem solution

In this HackerEarth Tree and K-Tuple problem solution You have a Tree with N nodes rooted at 1 where ith node is associated with a value Ai. Now consider the following definition:

K-Tuple: A tuple of K + 1 nodes (X1, X2, X3,…, Xk+1) is a K-Tuple if:

1.  Xi is an ancestor of Xi+1 for all 1<=i<=K
2.  Axi > Axi+1 for all 1<=i<=K

Calculate the number of K-Tuples in the tree.

#include <bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define ll long long
#define pll pair<ll, ll>
#define pii pair<int,int>
#define pb push_back
#define F first
#define S second
#define mod 1000000007
#define maxn 100005
#define boost ios::sync_with_stdio(false);cin.tie(0)
#define fr freopen("source.txt","r",stdin),freopen("output.txt","w",stdout)
#define SET(a,b) memset(a,b,sizeof(a))

ll a[maxn],sz[maxn],mark[maxn],t[maxn],temp[maxn],ans[maxn];
vector<int>g[maxn];

void update(int pos,ll val){
    val=(val+mod)%mod;
    for(int i=pos;i<maxn;i+=i&-i){
        t[i]+=val;
        if(t[i]>=mod)t[i]%=mod;
    }
}

ll read(int pos){
    ll res=0;
    for(int i=pos;i>0;i-=i&-i){
        res+=t[i];
        if(res>=mod)res-=mod;
    }
    return res;
}


void pre(int v,int p){
    sz[v]=1;
    for(int u:g[v]){
        if(u-p){
            pre(u,v);
            sz[v]+=sz[u];
        }
    }
}

void add(int v,int p){
    update(a[v],ans[v]);
    for(int u:g[v]){
        if(u-p&&mark[u]!=1){
            add(u,v);
        }
    }
}

void rem(int v,int p){
    update(a[v],-ans[v]);
    for(int u:g[v]){
        if(u-p&&mark[u]!=1){
            rem(u,v);
        }
    }
}

void dfs(int v,int p,int keep){
    int mx=-1;
    int bigchild=-1;
    for(int u:g[v]){
        if(u-p){
            if(mx<sz[u]){
                mx=sz[u];
                bigchild=u;
            }
        }
    }

    for(int u:g[v]){
        if(u-p && u!=bigchild){
            dfs(u,v,0);
        }
    }

    if(bigchild!=-1){
        dfs(bigchild,v,1);
        mark[bigchild]=1;
    }

    add(v,p);
    temp[v]=read(a[v]-1);

    if(bigchild!=-1){
        mark[bigchild]=0;
    }

    if(keep==0){
        rem(v,p);
    }
}

int main(){
    boost;
    
    int n,k;
    cin>>n>>k;

    rep(i,1,n)cin>>a[i];

    rep(i,2,n){
        int u,v;
        cin>>u>>v;
        g[v].pb(u);
        g[u].pb(v);
    }

    pre(1,0);

    rep(i,1,n)ans[i]=1;

    rep(i,1,k){
        dfs(1,0,0);
        rep(j,1,n)ans[j]=temp[j];
    }

    ll sum=0;
    rep(i,1,n)sum=(sum+ans[i])%mod;
    cout<<sum;

    return 0;
}