Link to the problem: https://practice.geeksforgeeks.org/contest/job-a-thon-for-internship/problems/#↵
~~~~~↵
ll dfs1(int src,int par,vector<int> val,ll dp1[],ll dp2[],int k,vector<vector<int>> adj){↵
if(dp1[src]!=-1)return dp1[src];↵
if(val[src]%k==0){↵
ll ans=0;↵
for(auto e:adj[src]){↵
if(e!=par){↵
ans+=dfs1(e,src,val,dp1,dp2,k,adj);↵
}↵
}↵
↵
dp1[src]=ans+1;↵
dp2[src]=0;↵
for(auto e:adj[src]){↵
if(e!=par)↵
dp2[src]+=((ans-dp1[e])*dp1[e]);↵
}↵
dp2[src]/=2;↵
dp2[src]+=ans;↵
// cout << src << " " << dp2[src] << " " << ans << endl;↵
return dp1[src];↵
}↵
else{↵
dp1[src]=0;↵
dp2[src]=0;↵
↵
for(auto e:adj[src]){↵
if(e!=par){↵
dfs1(e,src,val,dp1,dp2,k,adj);↵
}↵
}↵
return dp1[src];↵
}↵
}↵
↵
long long countPairs(int N, vector<vector<int>> Edges, vector<int> Val,int K)↵
{↵
vector<vector<int>> adj(N+5);↵
ll dp1[N+1];↵
memset(dp1,-1,sizeof dp1);↵
for(auto e: Edges){↵
int u=e[0];↵
int v=e[1];↵
adj[u].push_back(v);↵
adj[v].push_back(u);↵
}↵
ll dp2[N+1];↵
memset(dp2,0,sizeof dp2);↵
dfs1(0,-1,Val,dp1,dp2,K,adj);↵
↵
int str=0;↵
//cout << ans << endl;↵
↵
for(int i=0;i<N;i++){↵
str+=dp2[i];↵
}↵
return str;↵
}↵
~~~~~↵
↵
I have been trying to solve it but my solution either gives TLE or wrong answer and for this specific solution its TLE, Can anyone please find the error in my code and tell me the right approach.↵
Thanks and regards↵
ll dfs1(int src,int par,vector<int> val,ll dp1[],ll dp2[],int k,vector<vector<int>> adj){↵
if(dp1[src]!=-1)return dp1[src];↵
if(val[src]%k==0){↵
ll ans=0;↵
for(auto e:adj[src]){↵
if(e!=par){↵
ans+=dfs1(e,src,val,dp1,dp2,k,adj);↵
}↵
}↵
↵
dp1[src]=ans+1;↵
dp2[src]=0;↵
for(auto e:adj[src]){↵
if(e!=par)↵
dp2[src]+=((ans-dp1[e])*dp1[e]);↵
}↵
dp2[src]/=2;↵
dp2[src]+=ans;↵
// cout << src << " " << dp2[src] << " " << ans << endl;↵
return dp1[src];↵
}↵
else{↵
dp1[src]=0;↵
dp2[src]=0;↵
↵
for(auto e:adj[src]){↵
if(e!=par){↵
dfs1(e,src,val,dp1,dp2,k,adj);↵
}↵
}↵
return dp1[src];↵
}↵
}↵
↵
long long countPairs(int N, vector<vector<int>> Edges, vector<int> Val,int K)↵
{↵
vector<vector<int>> adj(N+5);↵
ll dp1[N+1];↵
memset(dp1,-1,sizeof dp1);↵
for(auto e: Edges){↵
int u=e[0];↵
int v=e[1];↵
adj[u].push_back(v);↵
adj[v].push_back(u);↵
}↵
ll dp2[N+1];↵
memset(dp2,0,sizeof dp2);↵
dfs1(0,-1,Val,dp1,dp2,K,adj);↵
↵
int str=0;↵
//cout << ans << endl;↵
↵
for(int i=0;i<N;i++){↵
str+=dp2[i];↵
}↵
return str;↵
}↵
↵
I have been trying to solve it but my solution either gives TLE or wrong answer and for this specific solution its TLE, Can anyone please find the error in my code and tell me the right approach.↵
Thanks and regards↵
