可撤销并查集模板

CF891C Envy

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#define inf 0x7fffffff
#define ll long long
//#define int long long
//#define double long double
#define re int
#define void inline void
#define eps 1e-5
//#define mod 1e9+7
//#define ls(p) p<<1
//#define rs(p) p<<1|1
//#define pi acos(-1.0)
#define pb push_back
#define P pair < int , int >
#define mk make_pair
#define fi first
#define se second
//#define unordered_map map
//#define __int128 long long
using namespace std;
using namespace __gnu_pbds;
const int M=1e5+5;
const int mod=1e9+7;
const int N=2e6+5;
vector < P > ee[N],q[N];
pair < P , int > e[N];
int ans[N];
int n,m;
namespace DSU
{
	int fa[N],rk[N];
	void init(int n){for(re i=1;i<=n;i++)  rk[i]=1,fa[i]=i;}
	int get(int x){return x==fa[x]?x:get(fa[x]);}
	void merge(int x,int y)
	{
		if(rk[x]==rk[y])  rk[x]++;
		if(rk[x]<rk[y])  swap(x,y);
		fa[y]=x;
	}
}
void undo(int w,int l,int r)
{
	int id=q[w][l].fi;
	vector < pair < P , P > > st;
	for(re i=l;i<r;i++)
	{
		int x=q[w][i].se;
		int u=e[x].fi.fi,v=e[x].fi.se;
		int fx=DSU::get(u),fy=DSU::get(v);
		if(fx==fy)  ans[id]=1;
		else
		{
			if(DSU::rk[fx]<DSU::rk[fy])  swap(fx,fy);
			st.pb(mk(mk(fy,DSU::fa[fy]),mk(fx,DSU::fa[fx])));
			if(DSU::rk[fx]==DSU::rk[fy])  DSU::rk[fx]++;
			DSU::fa[fy]=fx;
		}
	}
	for(re i=st.size()-1;i>=0;i--)
	{
		DSU::fa[st[i].fi.fi]=st[i].fi.se;
		DSU::rk[st[i].se.fi]=st[i].se.se;
	}
}
void solve()
{
	int Q;
	cin>>n>>m;
	DSU::init(n);
	for(re i=1;i<=m;i++)
	{
		int x,y,z;
		scanf("%d%d%d",&x,&y,&z);
		e[i]=mk(mk(x,y),z);
		ee[z].pb(mk(x,y));
	}
	cin>>Q;
	for(re i=1;i<=Q;i++)
	{
		int k;
		scanf("%d",&k);
		for(re j=1;j<=k;j++)
		{
			int x;
			scanf("%d",&x);
			q[e[x].se].pb(mk(i,x));
		}
	}
	for(re w=1;w<=505050;w++)
	{
		for(re l=0,r;l<q[w].size();l=r)
		{
			for(r=l;r<q[w].size()&&q[w][l].fi==q[w][r].fi;r++);
			undo(w,l,r); 
		}
		for(re i=0;i<ee[w].size();i++)
		{
			int x=ee[w][i].fi;
			int y=ee[w][i].se;
			int fx=DSU::get(x),fy=DSU::get(y);
			if(fx!=fy)  DSU::merge(fx,fy);
		}
	}
	for(re i=1;i<=Q;i++)  puts(ans[i]?"NO":"YES");
}
signed main()
{
//	fflush(stdout);
//	srand(102321547);
//	freopen("D.in", "r", stdin);
//	freopen("Dout.txt", "w", stdout);
//	freopen("9.out", "w", stdout);
    int T=1;
//    cin>>T;
    for(re index=1;index<=T;index++)
    {
//        printf("Case #%d:\n",index);
        solve();
//        puts("");
    }
    return 0;
}
/*



*/