Submission #927092

#	Time	Username	Problem	Language	Result	Execution time	Memory
927092		vjudge1	Paths (BOI18_paths)	C++17	0 / 100	141 ms	101968 KiB

paths

#include <bits/stdc++.h>

using namespace std;

#define all(a) a.begin(),a.end()
#define pb push_back
#define vt vector
#define endl '\n'
typedef long long ll;

const ll mod=1e9+7;
const ll inf=1e9+7;
const int N=1e6+4;

int n,m,k,a[N],used[N];
ll res=0;
vt<int>g[N][4];

void dfs(int v,int pr){
	used[v]=1;
	for(int i=1; i<=3; ++i) for(int to:g[v][i]) if(!used[to]) dfs(to,v);
	set<int>st;
	st.insert(1);
	st.insert(2);
	st.insert(3);
	st.erase(a[v]);
	auto bb=st.begin();
	auto ee=st.rbegin();
	res+=g[v][*ee].size()*g[v][*bb].size();
}

void dfs2(int v,int pr){
	used[v]=1;
	for(int i=1; i<=3; ++i) for(int to:g[v][i]) if(!used[to]) dfs2(to,v);
	set<int>st;
	st.insert(1);
	st.insert(2);
	st.insert(3);
	st.erase(a[v]);
	auto bb=st.begin();
	auto ee=st.rbegin();
	//cout<<*bb<<' '<<g[v][*bb].size()<<endl;
	res+=g[v][*bb].size()*2;
	res+=g[v][*ee].size()*2;
	if(pr!=-1 && a[v]!=a[pr]) res-=2;
	//cout<<v<<' '<<res<<endl;
}

void solve(){
	cin>>n>>m>>k;
	set<int>uni;
	for(int i=1; i<=n; ++i) {
		cin>>a[i];
	//	uni.insert(a[i]);
	}
	for(int i=1; i<=m; ++i){
		int u,v;
		cin>>u>>v;
		g[u][a[v]].pb(v);
		g[v][a[u]].pb(u);
	}
	if(k==1){
		cout<<0<<endl;
		return;
	}
	if(k==2){
		for(int i=1; i<=n; ++i) if(!used[i]) dfs2(i,-1);
		cout<<res<<endl;
		return;
	}
	if(k==3){
		for(int i=1; i<=n; ++i) if(!used[i]) dfs(i,-1);
		res*=2;
		fill(used+1,used+1+n,0);
		//cout<<res<<endl;
		for(int i=1; i<=n; ++i) if(!used[i]) dfs2(i,-1);
		cout<<res<<endl;
	}
}
			
int main(){
	int tt=1;
	//cin>>tt;
	while(tt--) {
		solve();
	}
}

• Subtask 1: 0 / 23
• Subtask 2: 0 / 20
• Subtask 3: 0 / 27
• Subtask 4: 0 / 30