kmjp's blog

競技プログラミング参加記です

AtCoder AGC #016 : B - Colorful Hats

これと今日のARCで赤色に戻れました。
http://agc016.contest.atcoder.jp/tasks/agc016_b

問題

N匹の猫がおり、それぞれある色の帽子をかぶっている。
数列Aが与えられる。
A[i]は、i匹目以外の猫の帽子の色数の総数である。
条件を満たすような帽子の色の組み合わせは存在するか判定せよ。

解法

N匹全部の帽子の色数をCとすると、A[i]はCかC-1のいずれかである。
よってAは最大で差が1まででなければならない。

  • A[i]がすべて同じ場合
    • これは以下の2通りが考えられる。
      • C=N-1で、N匹が皆異なる色の帽子をしている。
      • N≧2*Cであり、各色の帽子の猫が2匹以上いる。
  • A[i]がCまたはC-1である場合
    • 小さい方の(C-1)と答えた猫は、その猫しか持たない色の帽子を持っていることになる。
    • 逆に、Cと答えた猫はその猫と同じ色の帽子がそれ以外にある、すなわち全体で2匹以上いることになる。

上記条件をもとに不等式で解いてもよいが、条件漏れが怖かったので自分は後者は実際に色の割り当てを行って検証した。

int N;
map<int,int> M;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>x, M[x]++;
	if(M.size()>2) {
		return _P("No\n");
	}
	else if(M.size()==1) {
		x=M.begin()->first;
		
		if(x==N-1) return _P("Yes\n");
		if(x*2<=N) return _P("Yes\n");
		return _P("No\n");
	}
	else {
		auto X=*M.begin();
		auto Y=*M.rbegin();
		map<int,int> A;
		FOR(i,N) {
			if(i<X.second) {
				A[i]=1;
				if(A.size()>X.first) return _P("No\n");
			}
			else {
				x=X.second + (i-X.second)%(Y.first-X.second);
				A[x]++;
			}
		}
		map<int,int> B;
		FORR(r,A) {
			if(r.second==1) B[A.size()-1]++;
			else B[A.size()]+=r.second;
		}
		if(M==B) {
			_P("Yes\n");
		}
		else {
			_P("No\n");
		}
	}
}

まとめ

こういう場合分け問題はケース漏れしそうで怖い。