kmjp's blog

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

yukicoder : No.1017 Reiwa Sequence

本番出てたらさっと思いつけたかなぁ…。
https://yukicoder.me/problems/no/1017

問題

N要素の正整数列Aが与えられる。
各値の上限(=M)を150000とする。
それぞれに重み1,0,-1を付与して総和を取り、重みに非ゼロ要素が1個以上ある状態で、総和を0にしたい。
可能であれば重みの例を求めよ。

解法

2^k>Mkとなるkを考える。重みを0か1しか取れない場合、Aの先頭要素k個で作れる重み付きの総和はk~Mkの(M(k-1)+1)通りだが、重みの取り方は2^k通りある。
ということは鳩ノ巣原理より、同じ重みを取る組み合わせが2つ以上ある場合がある。
その重みの差を取れば解となる。

実際k=24で上記式を満たすので、先頭24要素の重みを適切にとればよいことになる。
おそらく上の通りに2^24通り試してもよいし、以下では先頭24要素を半分全列挙でそれぞれ3^12通りずつ求め、和が0になるケースを求めている。

int N;
int A[151515];

map<int,vector<int>> V;
vector<int> R;

void dfs(int cur,int sum) {
	if(R.size()==12) {
		V[sum]=R;
		return;
	}
	
	R.push_back(0);
	dfs(cur+1,sum);
	R.back()=A[cur];
	dfs(cur+1,sum+A[cur]);
	R.back()=-A[cur];
	dfs(cur+1,sum-A[cur]);
	R.pop_back();
}

void dfs2(int cur,int sum) {
	if(R.size()==12) {
		
		if(V.count(-sum)) {
			vector<int> R2=V[-sum];
			int ok=0;
			FORR(r,R) if(r!=0) ok++;
			FORR(r,R2) if(r!=0) ok++;
			if(ok) {
				int i;
				cout<<"Yes"<<endl;
				FOR(i,N) {
					if(i<12) cout<<R2[i]<<" ";
					else if(i<24) cout<<R[i-12]<<" ";
					else cout<<"0 ";
				}
				exit(0);
				
			}
			
		}
		return;
	}
	
	R.push_back(0);
	dfs2(cur+1,sum);
	R.back()=A[cur];
	dfs2(cur+1,sum+A[cur]);
	R.back()=-A[cur];
	dfs2(cur+1,sum-A[cur]);
	R.pop_back();
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) cin>>A[i];
	dfs(0,0);
	dfs2(12,0);
	cout<<"No"<<endl;
}

まとめ

鳩ノ巣原理から重みの差を取るのは今これ書いてて思いついた。
半分全列挙よりそっちが簡単だな。