本番出てたらさっと思いつけたかなぁ…。
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; }
まとめ
鳩ノ巣原理から重みの差を取るのは今これ書いてて思いついた。
半分全列挙よりそっちが簡単だな。