Submission #211860

#	Time	Username	Problem	Language	Result	Execution time	Memory
211860		zscoder	Making Friends on Joitter is Fun (JOI20_joitter2)	C++17	100 / 100	1311 ms	105436 KiB

joitter2

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;
 
#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key
 
typedef long long ll;
typedef pair<ll,ll> ii;
typedef vector<ll> vi;
typedef long double ld; 
typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> pbds;

ll global=0;

ll c2(ll x)
{
	return (x*(x-1));
}

set<ll> adj[122222]; //stores edges u->rt(v)
set<ll> radj[122222];
set<ll> a2[122222]; //stores edges rt(u)->rt(v)
set<ll> ra2[122222];
ll globaladj=0;
vector<ii> mergeq;

void transfer(ll r, ll r2);

struct DSU
{
	ll S;
	
	struct node
	{
		ll p;
		vi vec;
	};
	vector<node> dsu;
	
	DSU(ll n)
	{
		S = n;
		for(ll i = 0; i < n; i++)
		{
			node tmp;
			tmp.p = i; tmp.vec.pb(i);
			dsu.pb(tmp);
		}
	}
	
	void reset(ll n)
	{
		dsu.clear();
		S = n;
		for(ll i = 0; i < n; i++)
		{
			node tmp;
			tmp.p = i; tmp.vec.pb(i);
			dsu.pb(tmp);
		}
	}
	
	ll rt(ll u)
	{
		if(dsu[u].p == u) return u;
		dsu[u].p = rt(dsu[u].p);
		return dsu[u].p;
	}
	
	void merge(ll u, ll v)
	{
		u = rt(u); v = rt(v);
		if(u == v) return ;
		if(dsu[u].vec.size()<dsu[v].vec.size()) swap(u, v);
		dsu[v].p = u;
		global-=c2(dsu[u].vec.size());
		global-=c2(dsu[v].vec.size());
		ll sumsiz = dsu[u].vec.size()+dsu[v].vec.size();
		transfer(v,u);
		for(ll x:dsu[v].vec)
		{
			//we need to clear adj[x]!
			if(adj[x].find(u)!=adj[x].end())
			{
				globaladj-=sumsiz;
				adj[x].erase(u);
				radj[u].erase(x);
			}
			dsu[u].vec.pb(x);
		}
		dsu[v].vec.clear();
		global+=c2(dsu[u].vec.size());
	}
	
	bool sameset(ll u, ll v)
	{
		if(rt(u) == rt(v)) return true;
		return false;
	}
	
	ll getstat(ll u)
	{
		return dsu[rt(u)].vec.size();
	}
};

set<ll> naive[222222];
DSU dsu(1);
int main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	const ll DEBUG=0;
	for(ll cc=1;;cc++)
	{
		ll n,m; 
		if(!DEBUG) cin>>n>>m;
		else
		{
			n=rand()%7+2;
			m=rand()%(n*(n-1))+1;
		}
		mergeq.clear();
		vector<ii> E;
		if(DEBUG)
		{
			for(ll i=0;i<n;i++)
			{
				for(ll j=0;j<n;j++)
				{
					if(i!=j) E.pb({i,j});
				}
			}
			random_shuffle(E.begin(),E.end());
		}
		for(ll i=0;i<n;i++)
		{
			adj[i].clear();radj[i].clear();
			a2[i].clear(); ra2[i].clear();
			naive[i].clear();
		}
		globaladj=0;
		global=0;
		dsu.reset(n);
		vi res_naive;
		vi res_real;
		vector<ii> edges;
		for(ll i=0;i<m;i++)
		{
			ll u,v; 
			if(!DEBUG) {cin>>u>>v; u--; v--;} 
			else {u=E[i].fi; v=E[i].se;}
			edges.pb({u,v});
			if(DEBUG)
			{
				naive[u].insert(v);
				for(ll z=0;z<n;z++)
				{
					for(ll i=0;i<n;i++)
					{
						for(ll v:naive[i])
						{
							for(ll x:naive[v])
							{
								if(x==i) continue;
								if(naive[x].find(v)!=naive[x].end())
								{
									naive[i].insert(x);
								}
							}
						}
					}
				}
				ll sum=0;
				for(ll i=0;i<n;i++)
				{
					sum+=naive[i].size();
				}
				res_naive.pb(sum);
				//cerr<<"REAL SUM: "<<sum<<'\n';
			}
			
			ll r1 = dsu.rt(u); ll r2 = dsu.rt(v);
			//cerr<<"INSERT "<<r1<<' '<<r2<<'\n';
			if(a2[r2].find(r1)!=a2[r2].end())
			{
				//merge r1 and r2
				mergeq.pb({r1,r2});
				//dsu.merge(r1,r2);
				//cerr<<i<<' '<<"MERGE: "<<r1<<' '<<r2<<'\n';
				//cerr<<dsu.rt(r1)<<'\n';
			}
			while(!mergeq.empty()) 
			{
				ii tmp = mergeq.back(); mergeq.pop_back();
				dsu.merge(tmp.fi,tmp.se);
			}
			r1=dsu.rt(r1); r2=dsu.rt(r2);
			if(r1!=r2) 
			{
				if(adj[u].find(r2)==adj[u].end()) globaladj+=dsu.getstat(r2);
				adj[u].insert(r2);
				radj[r2].insert(u);
				a2[r1].insert(r2);
				ra2[r2].insert(r1);
			}
			ll ans=0;
			/*
			for(ll i=0;i<n;i++)
			{
				for(ll v:adj[i])
				{
					ll r1 = dsu.rt(i); ll r2 = dsu.rt(v);
					if(r1==r2) continue;
					//cerr<<"ADJ: "<<i<<' '<<v<<'\n';
					ans+=dsu.getstat(r2);
				}
			}
			*/
			ans+=global;
			ans+=globaladj;
			res_real.pb(ans);
			//cout<<ans<<'\n';
		}
		if(!DEBUG)
		{
			for(ll x:res_real)
			{
				cout<<x<<'\n';
			}
			return 0;
		}
		if(res_real!=res_naive)
		{
			for(ll i=0;i<m;i++)
			{
				cerr<<res_real[i]<<' '<<res_naive[i]<<'\n';
			}
			freopen("joitter-tc.out","w",stdout);
			cout<<n<<' '<<m<<'\n';
			for(ii e:edges)
			{
				cout<<e.fi+1<<' '<<e.se+1<<'\n';
			}
			return 0;
		}
		cerr<<"Case #"<<cc<<" complete\n";
	}
}

void transfer(ll r, ll r2)
{
	ll oldsiz = dsu.dsu[r].vec.size();
	ll newsiz = dsu.dsu[r].vec.size()+dsu.dsu[r2].vec.size();
	globaladj+=ll(oldsiz)*ll(radj[r2].size()); //add these to value, then now we normalize by removing stuff
	for(ll x:radj[r])
	{
		assert(adj[x].find(r)!=adj[x].end());
		//cerr<<"REMOVE: "<<x<<' '<<r<<'\n';
		globaladj-=oldsiz;
		adj[x].erase(r);
		if(dsu.rt(x)!=r2)
		{
			if(adj[x].find(r2)==adj[x].end())
			{
				globaladj+=newsiz;			
				adj[x].insert(r2);
			}
			radj[r2].insert(x);
		}
	}
	radj[r].clear();
	for(ll x:ra2[r])
	{
		a2[x].erase(r);
		if(x!=r2)
		{
			a2[x].insert(r2);
			ra2[r2].insert(x);
			if(a2[r2].find(x)!=a2[r2].end())
			{
				mergeq.pb({x,r2});
			}
		}
	}
	for(ll x:a2[r])
	{
		ra2[x].erase(r);
		if(x!=r2)
		{
			a2[r2].insert(x);
			ra2[x].insert(r2);
			if(a2[x].find(r2)!=a2[x].end())
			{
				mergeq.pb({x,r2});
			}
		}
	}
	a2[r].clear();
	ra2[r].clear();
}

Compilation message (stderr)

joitter2.cpp: In function 'int main()':
joitter2.cpp:246:11: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)', declared with attribute warn_unused_result [-Wunused-result]
    freopen("joitter-tc.out","w",stdout);
    ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 1 / 1
• Subtask 2: 16 / 16
• Subtask 3: 83 / 83