HackerEarth GCD Path on a Tree problem solution

In this HackerEarth GCD Path on a Tree problem solution Given an Integer K and a tree of N vertices where each vertex consists of a value Vi associated to it. Find the longest path in the tree which satisfies the following conditions

The number of vertices in the path should be a multiple of K.

Let’s say there are L X K vertices in the path and Xi be the value associated with ith vertex in that path, then gcd(X(iXK+1), X(iXK+2), X(iXK+3), … X(iXK+K)) > 1 for all i belongs to [0,L – 1] where gcd(x1,x2…xn) is the Greatest Common Divisor of the numbers x1,x2…xn.

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef double D;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef pair<int,int> PII;
typedef pair<ll,ll> PLL;

#define rd(x) scanf("%d",&x)
#define rd2(x,y) scanf("%d %d",&x,&y)
#define rl(x) scanf("%lld",&x)
#define rl2(x,y) scanf("%lld %lld",&x,&y)
#define wd(x) printf("%d" ,x)
#define wd2(x,y) printf("%d %d",x,y)
#define wl(x) printf("%lld",x)
#define wl2(x,y) printf("%lld %lld",x,y)
#define PC(x) putchar(x)
#define GC() getchar()
#define NL printf("n")
#define SP printf(" ")

#define F first
#define S second
#define MP make_pair
#define PB push_back

#define fr(i,s,e) for(int i=s;i<e;i++)
#define frr(i,s,e) for(int i=s-1;i>=e;i--)
#define frv(i,a) for(int i = 0;i<(int)a.size();i++)
#define frvr(i,a) for(int i = a.size()-1;i>=0;i--)
#define tr(i,a) for(typeof(a.begin()) i = a.begin(); i != a.end();i++)

#ifdef LOCAL
#define E1(a) cout<<#a<<":" <<a<<endl;
#define E2(a,b) cout<<#a<<":" <<a<<" " <<#b<<":" <<b<<endl;
#define E3(a,b,c) cout<<#a<<":" <<a<<" " <<#b<<":" <<b<<" "<<#c<<":"<<c<<endl;
#define E4(a,b,c,d) cout<<#a<<":" <<a<<" " <<#b<<":" <<b<<" "<<#c<<":"<<c<< " "<<#d<< ":"<<d<<endl;
#define INP freopen("input","r",stdin);
#define OUT freopen("output","w",stdout);
#else
#define E1(a)
#define E2(a,b)
#define E3(a,b,c)
#define E4(a,b,c,d)
#define INP
#define OUT
#endif

#define mod 1000000007
#define maxn 100009
#define maxr 1000009




int n,k,ans,val[maxn],mx[maxn];
VI pf[maxr],al[maxn];
map<int,VI > dp[maxn];

void primeFactors(vector<int> v[],int n){
    for(int i=2;i<=n;i++){
        if(v[i].empty()){
            for(int j=i;j<=n;j+=i){
                v[j].push_back(i);
            }
        }
    }
}

bool inline check(int v,int p,int m){
    if(dp[v].find(p)==dp[v].end()||dp[v][p][m]==-1)return 0;
    return 1;
}

void calc(int v,int p){
    map<int,VI> bst_mp;
    int bst = 0;

    VI &factors = pf[val[v]];
    mx[v] = 0;

    frv(i,factors){
        int prime = factors[i];
        dp[v][prime] = vector<int>(k+1,-1);
        bst_mp[prime] = vector<int>(k+1,-1);
        dp[v][prime][1] = 1;
        if(k==1){dp[v][prime][k] = mx[v] = 1;}
    }

    frv(i,al[v]){
        int u = al[v][i];
        if(u==p)continue;

        frv(j,factors){
            int prime = factors[j];

            dp[v][prime][1] = max(dp[v][prime][1],mx[u]+1);
            for(int m=2;m<=k;m++){
                if(!check(u,prime,m-1))continue;
                dp[v][prime][m] = max(dp[v][prime][m],dp[u][prime][m-1]+1);

            }
            mx[v] = max(mx[v],dp[v][prime][k]);
        }

        if(k==1){ans = max(ans,bst+1+mx[u]);}
        else{
            frv(j,factors){
                int prime = factors[j];

                if(bst_mp[prime][k-1]!=-1){ans = max(ans,bst_mp[prime][k-1]+1+mx[u]);}
                if(check(u,prime,k-1)){ans = max(ans,bst+1+dp[u][prime][k-1]);}

                for(int m=1;m<=k-2;m++){
                    if(bst_mp[prime][k-m-1]!=-1&&check(u,prime,m)){
                        ans = max(ans,bst_mp[prime][k-m-1]+1+dp[u][prime][m]);
                    }
                }
            }

            frv(j,factors){
                int prime = factors[j];
                for(int m=1;m<=k;m++){
                    if(check(u,prime,m))bst_mp[prime][m] = max(bst_mp[prime][m],dp[u][prime][m]);
                }
            }
        }
        bst = max(bst,mx[u]);
    }
    ans = max(ans,mx[v]);
}

void dfs(int v,int p){
    frv(i,al[v]){
        int u = al[v][i];
        if(u==p)continue;
        dfs(u,v);
    }
    calc(v,p);
}


int main(){

    for(int i=0;i<maxr;i++)pf[i].clear();
    primeFactors(pf,maxr-1);

    cin>>n>>k;

    for(int i=0;i<=n;i++)al[i].clear();


    for(int i=0;i<n-1;i++){
        int u,v;
        rd2(u,v);
        al[u].PB(v);
        al[v].PB(u);
    }

    for(int i=1;i<=n;i++)rd(val[i]);

    ans = 0;
    dfs(1,0);
    cout<<ans<<endl;
}