Submission #958236

---

Source Code Expand

#include <bits/stdc++.h>
#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<int> vi;
typedef long double ld; 
typedef tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update> pbds;
typedef set<int>::iterator sit;
typedef map<int,int>::iterator mit;
typedef vector<int>::iterator vit;
 
vi adj[100001];
int out[100001];
int in[100001];
vi type[100001];
int color[100001];
bool used2 = false;
int cnt = 0;
int c[3];
ll mm = 0;
ll tote = 0;
bool dfs(int u, int col)
{
	used2|=(in[u]&&out[u]);
	cnt++; c[col]++;
	color[u] = col;
	bool res = 1;
	mm += out[u];
	tote++;
	for(int i = 0; i < adj[u].size(); i++)
	{
		int v = adj[u][i];
		int nw = (col+type[u][i])%3;
		if(color[v] == -1)
		{
			res&=dfs(v, nw);
		}
		else
		{
			if(color[v] != nw) res=false;
		}
	}
	return res;
}
 
int main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	int n, m; cin>>n>>m;
	for(int i = 0; i < m; i++)
	{
		int u, v; cin>>u>>v; u--; v--;
		adj[u].pb(v); adj[v].pb(u);
		type[u].pb(1); type[v].pb(2);
		out[u]++; in[v]++;
	}
	memset(color,-1,sizeof(color));
	ll res = 0;
	for(int i = 0; i < n; i++)
	{
		if(color[i] == -1)
		{
			used2 = false; cnt = 0;
			memset(c,0,sizeof(c));
			mm = 0; tote = 0;
			bool ans=dfs(i, 0);
			//cerr<<c[0]<<' '<<c[1]<<' '<<c[2]<<' '<<mm<<'\n';
			if(ans)
			{
				if(used2)
				{
					for(int j = 0; j < 3; j++)
					{
						res += ll(c[j])*ll(c[(j+1)%3]);
					}
				}
				else
				{
					res += mm;
				}
			}
			else
			{	
				res += ll(tote)*ll(tote);
			}
		}
	}
	cout<<res<<'\n';
}

Submission Info

Judge Result