Codeforces Round #176 (Div.1) E. Ladies' Shop
解法
FFT.
まず,答えの集合から生成できる任意の重さは,ある2個以下のカバンを使って(同じカバンを2つでもいい)作ることができる.
なぜなら,もしある重さを p1 + p2 + p3 で生成したとする.もちろん p3 は与えられたカバンに含まれているはずである.また,p1 + p2 も同様に,与えられたカバンに含まれているはずである.そうでなければ集合の選び方に矛盾する.
したがって,問題を言い換えると,
「与えられたかばんの重さからちょうど2つ(同じものを選んで良い)選んで作れる重さの集合が,与えられたかばんの集合の部分集合であることを判定し,さらにその集合を生成する最小の元の個数を求めよ」
となる.
これはFFTで計算できる.第 a[i] 項を1としてそれ以外を0とする多項式 P をつくり,FFT で P^2 を求めれば良いからである.
また,答えの集合に含めなければならない重さは,P^2 のうち係数が0 かつ かばんにその重さが含まれているものである.
このために,Pの第0項は場合の数を求めるときとはうってかわって0にして置かなければならない.
ソースコード
// convolution は略 // data_type := complex<double>; int main() { int n, m; scanf("%d %d", &n, &m); vector<int> a(n); vector<data_type> p(m + 1); unordered_set<int> s; for(int i = 0; i < n; ++i) { scanf("%d", &a[i]); p[a[i]] = data_type(1, 0); s.insert(a[i]); } auto p2 = convolution(p, p); vector<int> ans; bool ok = true; for(int i = 0; i <= m; ++i) { const int c = real(p2[i]) + 0.1; if(c == 0) { if(s.count(i) == 1) ans.push_back(i); continue; } ok &= s.count(i); } if(ok) { printf("YES\n%d\n", (int)ans.size()); for(int i = 0; i < (int)ans.size(); ++i) { if(i != 0) printf(" "); printf("%d", ans[i]); } printf("\n"); } else { printf("NO\n"); } }
感想
入出力が重いので注意.
自分のFFTはお世辞にも早いとは言えないので,TLEしてしまった.scanf/printf すると通る.