codeforces 884C Bertown Subway

http://codeforces.com/problemset/problem/884/C

找环，再按每个环内元素的数量从大到小排序，把前2个大的求和，再将这个和与其他的数求平方和。

#include <bits/stdc++.h>
using namespace std;
int con=0;long long int  ans[111111]; bool vis[111111];    int ma[111111];
void dfs(int x,int num)
{
	if(vis[ma[x]]==1)
	{
		ans[con]=num-1;
		con++;
		return;
	}
	else {
		vis[ma[x]]=1;
		dfs(ma[x],num+1);
	}
}
int main()
{
   long long int n;

    
    
    while(cin>>n){
    	memset(ma,0,sizeof(ma));
    	memset(vis,0,sizeof(vis));
		memset(ans,0,sizeof(ans)); 
    	con=0;
    	int i;
    	int x,y;
    	for(i=1;i<=n;i++)
    	{
    		scanf("%d",&x);
    		ma[i]=x;
		}
		for(i=1;i<=n;i++)
		{
			if(vis[i]==0)
			{
				vis[i]=1;
				dfs(i,2);
			}
		}
		//for(i=0;i<con;i++)
		//cout<<ans[i]<<" ";
		//cout<<endl;
		if(con<=2){
			long long int a=n*n;
			cout<<a<<endl;
		}
		else{
			sort(ans,ans+con);
			long long int a=0;
			a+=(ans[con-1]+ans[con-2])*(ans[con-1]+ans[con-2]);
			for(int j=0;j<=con-3;j++)
			{
				a+=ans[j]*ans[j];
			}
			cout<<a<<endl;
		}
	}
    return 0;
}