求助卡常
查看原帖
求助卡常
564225
dontwannacry楼主2024/3/9 19:01

思路是将每个连通块缩点之后预处理两个颜色的情况

目前看来正确性应该是没问题的,但复杂度可能假了

提交记录

#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;
}
2024/3/9 19:01
加载中...