I watched the video editorial of FLGZRO and tried to implement the solution. But I am getting TLE. Could someone please point out any optimization?
#include<bits/stdc++.h>
using namespace std;
#define ll long long
#define endl "\n"
const int M = 1e9+7;
vector<int> graph[100001];
ll in[100001], ans[100001], subtree[100001], a[100001];
map<ll,vector<int>> aval;
int timer;
void dfs(int node, int par){
ans[node] = 0;
subtree[node] = 0;
in[node] = timer++;
for(auto child:graph[node]){
if(child==par)continue;
dfs(child,node);
ans[node] += subtree[child];
subtree[node] += subtree[child];
ans[node] %= M;
subtree[node] %= M;
}
if(ans[node]==0){
ans[node] = 1;
subtree[node] = 1;
}
else{
for(auto it: aval[a[node]]){
if(in[node]<in[it]){
ans[node] += (M - ans[it]);
ans[node] %= M;
}
}
subtree[node] += ans[node];
subtree[node] %= M;
}
aval[a[node]].push_back(node);
}
int main(){
ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
int t; cin>>t;
while(t--){
int n; cin>>n;
timer = 1;
for(int i = 0; i <n -1; i++){
int u,v;
cin>>u>>v;
graph[u].push_back(v);
graph[v].push_back(u);
}
for(int i =1; i <=n; i++)cin>>a[i];
if(n==1){
cout<<1<<endl;
continue;
}
dfs(1,-1);
cout<<subtree[1]<<endl;
for(int i = 0; i <=n; i++){
graph[i].clear();
in[i] = 0;
}
aval.clear();
}
return 0;
}