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.

HackerEarth Tree and K-Tuple problem solution.

#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;
}

Second solution

#include <bits/stdc++.h>
using namespace std;
vector<int>g[100011];
int bt[100011][22];
const int mod = 1e9+7;
void update(int k,int ind,int val) {
    while(ind) {
        bt[ind][k] += val;
        if(bt[ind][k]<0) bt[ind][k]+=mod;
        if(bt[ind][k]>=mod) bt[ind][k]-=mod;
        ind-=(ind&-ind);
    }
}
int query(int k,int ind) {
    int ret = 0;
    while(ind<=100000) {
        ret += bt[ind][k];
        if(ret>=mod) ret-=mod;
        ind+=(ind&-ind);
    }
    return ret;
}

int ans = 0;
int N,k;
int dp[100011][22];
int A[100011];
void dfs(int v,int p) {
    dp[v][1] = 1;
    for(int i=2;i<=k+1;i++) {
        dp[v][i] = query(i-1,A[v]+1);
    }
    for(int i=1;i<=k+1;i++) {
        update(i,A[v],dp[v][i]);
    }
    ans += dp[v][k+1];
    if(ans >=mod) ans-=mod;

    for(auto x:g[v]) {
        if(x!=p) {
            dfs(x,v);
        }
    }
    for(int i=1;i<=k+1;i++) {
        update(i,A[v],-dp[v][i]);
    }
}
int main()
{
    int u,v;
    cin >> N >> k;
    for(int i=1;i<=N;i++) {
        cin >> A[i];
    }
    for(int i=0;i<N-1;i++) {
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1,0);
    cout << ans;
}