日立製作所 社会システム事業部 プログラミングコンテスト2020 : C - ThREE

Eをえいやで解いてしまったけど、おかげでどうにか赤復帰。
https://atcoder.jp/contests/hitachi2020/tasks/hitachi2020_c

問題

木を成すN頂点の無向グラフが与えられる。
頂点に1～Nの値を1つずつ振るとき、距離3の頂点間で値の和か積が3の倍数となるようにしたい。
振り方を一つ答えよ。

解法

1～Nを3で割った余りが0,1,2のいずれかで分類する。
許可されないのは、1または2で同じ側に分類される点が距離3の頂点間に割り振られることである。
一方余りが0の値はどこにでも置けることになる。

距離3をまじめに考えると大変なので、二部グラフを考える。
余りが1または2になる値が分けた両側にきてしまうと、条件に反してしまう可能性がある。
そこで、そのようなケースをなんとか避けよう。

• 二部グラフの両側の頂点数がいずれもN/3以上の場合
  • 片側に余りが1、反対側に余りが2のものをすべて置いてしまおう。残りは余りが0のものを置く。
• 上記以外の場合
  • この場合、二部グラフのいずれかは2N/3以上の頂点数になる。そこに余りが1と2のものを押し込んでしまおう。

int N;
vector<int> E[202020];
vector<int> C[3];
int R[202020];

vector<int> cand[2];

void dfs(int cur,int pre,int d) {
	cand[d].push_back(cur);
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur,d^1);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	for(i=1;i<=N;i++) C[i%3].push_back(i);
	
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,0,0);
	
	if(cand[0].size()>cand[1].size()) swap(cand[0],cand[1]);
	
	if(cand[0].size()<=C[0].size()) {
		while(cand[0].size()&&C[0].size()) {
			x=cand[0].back();
			y=C[0].back();
			cand[0].pop_back();
			C[0].pop_back();
			R[x]=y;
		}
	}
	else {
		assert(C[2].size()<=cand[0].size());
		assert(C[1].size()<=cand[1].size());
		while(cand[0].size()&&C[2].size()) {
			x=cand[0].back();
			y=C[2].back();
			cand[0].pop_back();
			C[2].pop_back();
			R[x]=y;
		}
		while(cand[1].size()&&C[1].size()) {
			x=cand[1].back();
			y=C[1].back();
			cand[1].pop_back();
			C[1].pop_back();
			R[x]=y;
		}
	}
	
	FORR(c,cand[0]) {
		FOR(i,3) if(C[i].size()) {
			R[c]=C[i].back();
			C[i].pop_back();
			break;
		}
	}
	FORR(c,cand[1]) {
		FOR(i,3) if(C[i].size()) {
			R[c]=C[i].back();
			C[i].pop_back();
			break;
		}
	}
	
	
	
	FOR(i,N) cout<<R[i]<<" ";
	
}

まとめ

最近CodeforcesのどこかのDiv1Cで、頂点を3つのグループに分けるとどこかN/3以上になることを用いた問題が解けなかった記憶があったので、こちらはすんなり思いつけた。