CW 模拟赛 T1.Soso 的并查集写挂了

Yorg / 2024-10-29 / 原文

题面

似乎有原题, 但是很偏
挂个 pdf
题面下载

算法

暴力

很显然, 只需要在并查集维护时稍微加上一点细节

#include <cstdio>
using namespace std;
int n,m,fa[500010],a[500010];
long long ans=0;
int find(int x){
  	ans+=a[x];
	ans%=998244353;
	if(fa[x]==x) return x;
	return find(fa[x]);
}
void merge(int x,int y){
	int tx=find(x);
  	int ty=find(y);
  	fa[ty]=tx;
}
int main(){
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		fa[i]=i;
	}
	for(int i=1,x,y;i<=m;i++){
		scanf("%d%d",&x,&y);
		merge(x,y);
	}
	printf("%lld\n",ans);
	return 0;
}

正解

观察到 TLE 的原因是每次查询使用了大量时间, 考虑优化

带权并查集

并查集的优点是能够通过路径压缩优化复杂度

维护边权

这是简单的, 只需要稍稍修改一下 $\rm{find}$ 函数即可 (其中 $d$ 存储到根的边权之和)

inline int find(int x) {
	if(fa[x] == x) return x;
	int root = find(fa[x]); //注意一下写法，先将find(fa[x])存放在root中，否则会出错 
	d[x] += d[fa[x]];
	return fa[x] = root;
}

维护点权

按照维护边权的方法写完之后发现会出问题
观察到本质上是因为每一次路径压缩都会重复计算最上层的根节点
于是想办法消除其的影响

代码(后补)

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int mod=998244353;
int n,m;
int a[500010];
pair<int,int> fa[500010];
int ans;
void init(){
	for(int i=1;i<=n;i++){
		fa[i]={i,a[i]};
	}
	return;
}
pair<int,int> find(int x){
	if(fa[x].first==x)return fa[x];
	auto t=find(fa[x].first);
	t.second+=fa[x].second;
	fa[x].second=t.second-fa[t.first].second; 
	fa[x].first=t.first;
	fa[x].second%=mod;
	t.second%=mod;
	return t;
}
signed main(){
	cin>>n>>m;
	for(int i=1;i<=n;i++){
		cin>>a[i];
	}
	init();
	while(m--){
		int x,y;
		cin>>x>>y;
		auto fx=find(x),fy=find(y);
		ans+=fx.second+fy.second;
		if(fx.first!=fy.first){
			fa[fy.first].first=fx.first;
		}
	}
	cout<<ans%mod;
	return 0; 
}

树上差分

观察到查询操作在最终的树上都是一条链, 考虑并查集 + 建树, 维护树上差分操作

#include <cstdio>
#include <vector>
#include <cstring>
#include <numeric>
#include <algorithm>
using namespace std;
#ifdef LOCAL
#define debug(...) fprintf(stderr,##__VA_ARGS__)
#else
#define debug(...) void(0)
#endif
typedef long long LL;
const int P=998244353;
vector<int> g[1<<19];
template<int N> struct dsu{
	int fa[N+10],cnt;
	explicit dsu(int n=N):cnt(n){iota(fa+1,fa+n+1,1);}
	int find(int x){return fa[x]==x?x:fa[x]=find(fa[x]);}
	void merge(int x,int y){if(x=find(x),y=find(y),x!=y) cnt--,fa[y]=x,g[x].push_back(y);}
};
int n,m,cnt[1<<19],a[1<<19];
dsu<1<<19> s;
LL ans=0;
void dfs(int u){
	for(int v:g[u]) dfs(v),cnt[u]+=cnt[v];
	ans=(ans+1ll*cnt[u]*a[u])%P;
}
int main(){
//	#ifdef LOCAL
//	 	freopen("input.in","r",stdin);
//	#endif
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++) scanf("%d",&a[i]);
	for(int i=1,u,v;i<=m;i++){
		scanf("%d%d",&u,&v);
		cnt[u]++,cnt[v]++;
		cnt[s.find(u)]--,cnt[s.find(v)]--;
		ans=(ans+a[s.find(u)]+a[s.find(v)])%P;
		s.merge(u,v);
	}
	for(int i=1;i<=n;i++) if(s.find(i)==i) dfs(i);
	printf("%lld\n",ans);
	return 0;
}