思路是将每个连通块缩点之后预处理两个颜色的情况
目前看来正确性应该是没问题的,但复杂度可能假了
#include<bits/stdc++.h>
using namespace std;
struct both{
int x,y;
bool operator < (const both other)const{
if(x == other.x) return y < other.y;
return x < other.x;
}
};
int vis[1000010],color2[1000010],si[1000010],cnt;
int color[1000010],in,ans[1000010];
vector<int>ne[1000010],bigne[1000010];
unordered_map<int,bool>to[1000010];
map<both,long long>add;
bool vis2[1000010];
int getblock(int now,int x){
vis[now] = x;
int res = 1;
for(int i = 0;i < ne[now].size();++i){
int v = ne[now][i];
if(vis[v]){
bigne[x].push_back(vis[v]);
bigne[vis[v]].push_back(x);
to[x][vis[v]] = 1;
to[vis[v]][x] = 1;
}
else if(color[v] == color[now])res += getblock(v,x);
}
return res;
}
int dfs2(int now,int x,int y,long long &sum){
vis2[now] = 1;
int res = si[now];
sum += si[now]*si[now];
for(int i = 0;i < bigne[now].size();++i){
int v = bigne[now][i];
if(vis2[v])continue;
if(color2[v] == x||color2[v] == y){
to[now][v] = to[v][now] = 0;
res += dfs2(v,x,y,sum);
}
}
return res;
}
int main(){
ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
int N,Q;cin >> N >> Q;
for(int i = 1;i <= N;++i){cin >> color[i];}
for(int i = 2;i <= N;++i){
cin >> in;
ne[in].push_back(i);
ne[i].push_back(in);
}
for(int i = 1;i <= N;++i){
if(!vis[i]){
++cnt;
si[cnt] = getblock(i,cnt);
color2[cnt] = color[i];
ans[color2[cnt]] += si[cnt]*si[cnt];
}
}
for(int i = 1;i <= N;++i){
for(int t = 0;t < bigne[i].size();++t){
if(to[i][bigne[i][t]]){
long long res,sum = 0;
res = dfs2(i,color2[i],color2[bigne[i][t]],sum);
add[{color2[i],color2[bigne[i][t]]}] = add[{color2[bigne[i][t]],color2[i]}] += res*res-sum;
}
}
}
while(Q--){
int x,y;cin >> x >> y;
//cout << ans[x]<<" "<<ans[y]<<" "<<add[{x,y}]<<"\n";
cout << ans[x]+ans[y]+add[{x,y}] <<"\n";
}
return 0;
}