AOJ 2390 AYBABTU

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2390

解法

感覚的には最大全域木的なことを基地が繋がりすぎないようにやると通りそう．
そして投げたら通った（終了）…だと困るので厳密じゃないけど一応それっぽいのだけ．

木なので基地でパスを区切ったものを考えれば，max(パスのなかの辺の重み)の小さい方から k - 1 個使えば嬉しい．
これは普通に辺のコストをでかいほうから使っていくクラスカルの要領でやっていき，加えた辺を負の値と考えて答えから省くようにすると同じことができる．
加えようとする辺が基地と基地を結ぶようなものであるときは，それ以上くっつかれると困る場合は無視，そうでない場合は使って上限をへらせばよい．

ソースコード

#include <bits/stdc++.h>
using namespace std;

class union_find {
public:
    union_find(int n) : par(n, -1) {}
    int root(int x) { return par[x] < 0 ? x : par[x] = root(par[x]); }
    void unite(int x, int y) {
        x = root(x), y = root(y);
        if(x == y) return;
        if(par[x] > par[y]) swap(x, y);
        par[x] += par[y];
        par[y] = x;
    }
private:
    vector<int> par;
};

struct edge {
    int u, v, cost;
    bool operator<(edge const& e) const {
        return cost < e.cost;
    }
};

int main() {
    int n, t, k;
    int cs = 1;
    while(cin >> n >> t >> k, n) {
        vector<edge> es(n - 1);
        int ans = 0;
        for(int i = 0; i < n - 1; ++i) {
            cin >> es[i].u >> es[i].v >> es[i].cost;
            es[i].u--; es[i].v--;
            ans += es[i].cost;
        }
        vector<bool> is_base(n);
        for(int i = 0; i < t; ++i) {
            int b; cin >> b;
            is_base[b - 1] = true;
        }
        sort(rbegin(es), rend(es));
        union_find uf(n);
        int can = t - k - 1;
        for(auto& e : es) {
            int u = uf.root(e.u), v = uf.root(e.v);
            if(is_base[u] && is_base[v]) {
                if(can > 0) {
                    can--;
                    ans -= e.cost;
                    uf.unite(u, v);
                }
            } else {
                uf.unite(u, v);
                ans -= e.cost;
                bool b1 = is_base[u], b2 = is_base[v];
                is_base[u] = is_base[u] | b2;
                is_base[v] = is_base[v] | b1;
            }
        }

        cout << "Case " << cs++ << ": " << ans << endl;
    }
}

感想

これはAOJ-ICPC700らしいんですが，体感300以下．厳密な証明が難しいから700？わからん．

AOJ 1342 Don't Burst the Balloon

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1342

解法

3次元幾何に見せかけた2次元幾何．
球の半径Rを決めた時，上から見て中心をどこにおけるかをイメージする．
すると針と球の条件は，ある円があってその外側に中心を置く，という感じになる．この円の半径は針の高さとRから簡単に求まる．
同様に壁との関係はある正方形があってその内部，になる．

こうして得られた正方形と円（針ごとに1つある）にたいして，適当に点を選んでその点が条件を満たすか考える．この点は，円と円の交点，円と正方形の交点，正方形の4隅を選べば十分．条件を満たすか確認するときは若干制約を緩くしておく（そうしないと周上にいる判定になってしまう）．

これができたら，あとは R で二分探索しておわり．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ld = long double;
using point = complex<ld>;

// 略

int main() {
    int n, w;
    while(cin >> n >> w, n) {
        vector<ld> x(n), y(n), h(n);
        for(int i = 0; i < n; ++i) {
            cin >> x[i] >> y[i] >> h[i];
        }

        auto check = [&](const ld R) {
            vector<circle> cs(n);
            for(int i = 0; i < n; ++i) {
                if(h[i] <= R) cs[i].r = sqrt(2 * R * h[i] - h[i] * h[i]);
                else          cs[i].r = R;
                cs[i].p = point(x[i], y[i]);
            }

            const ld margin = (w <= R ? sqrt(2 * R * w - w * w) : R);
            if(margin >= 50) return false;

            vector<segment> ss;
            ss.emplace_back(point(margin, margin), point(100 - margin, margin));
            ss.emplace_back(point(margin, margin), point(margin, 100 - margin));
            ss.emplace_back(point(100 - margin, margin), point(100 - margin, 100 - margin));
            ss.emplace_back(point(margin, 100 - margin), point(100 - margin, 100 - margin));

            vector<point> ps = {point(margin, margin), point(margin, 100 - margin),
                                point(100 - margin, margin), point(100 - margin, 100 - margin)};
            for(int i = 0; i < n; ++i) {
                for(int j = 0; j < 4; ++j) {
                    for(auto&& p : is_sc(cs[i], ss[j])) ps.push_back(p);
                }
                for(int j = i + 1; j < n; ++j) {
                    for(auto&& p : is_cc(cs[i], cs[j])) ps.push_back(p);
                }
            }

            bool ok = false;
            for(auto& p : ps) {
                bool b = true;
                for(auto& c : cs) {
                    b &= abs(p - c.p) >= c.r - eps;
                }
                b &= real(p) >= margin - eps && real(p) - eps <= 100 - margin;
                b &= imag(p) >= margin - eps && imag(p) - eps <= 100 - margin;

                ok |= b;
            }

            return ok;
        };

        ld lb = 0, ub = 1e9;
        for(int lp = 0; lp < 50; ++lp) {
            const ld mid = (ub + lb) / 2;
            if(check(mid)) lb = mid;
            else           ub = mid;
        }

        cout << fixed << setprecision(10) << lb << endl;
    }
}

感想

なんか前書いた時バグった記憶があるんだけど，今書いたら15分ぐらいでかけてよかった．

AOJ 2558 Medical Inspection

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2558

解法

最小点被覆は多項式時間で解けない．
この時点で，グラフにある良い性質を用いるか，頑張って枝刈りしかない．
しかし与えられたグラフは単純グラフであるだけで他にいい感じの性質がない．
よって，泣きながら枝刈りを書く．

枝刈りの戦略は以下のようにした．何が良いかよくわからないけど．

• 次数が k + 1 以上なら使わなければ死ぬので最初から使っておく（これは絶対必要）
• 探索するときは，次数の多い方から選んでいくほうが嬉しそう？（これがどれぐらい寄与してるかよくわからない）
• 注目している頂点を使わない場合は，隣接頂点をすべて使うのが確定する．なので，この時点で隣接頂点すべてを削除しておくこと．（自分はこれを忘れていてTLE喰らいまくった）
• 上記の戦略で探索していくなら，次数1のやつはスキップするのがよい．どうせ後でいい感じに消えるし，次数1だけのために貴重な k を使って探索空間を広げても嬉しくない．
• 探索過程で，(使える残りの数) * (余ってる頂点で最大の次数) < (残ってる辺の数) ならどうやっても無理なのでこれも枝刈り（これも多分必要だと思う）．

結局 0.05sec で通った．
なんか調べてみたら，O(1.3^k + kn) ぐらいで解く指数時間アルゴリズムがあるらしい．やばそう．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using pii = pair<int, int>;

int main() {
    int n, m, k;
    cin >> n >> m >> k;
    vector<set<int>> g(n);
    vector<int> deg(n);
    for(int i = 0; i < m; ++i) {
        int a, b; cin >> a >> b;
        g[a - 1].insert(b - 1);
        g[b - 1].insert(a - 1);
        deg[a - 1]++;
        deg[b - 1]++;
    }

    vector<bool> used(n);
    int now = 0;
    for(int i = 0; i < n; ++i) {
        if(deg[i] > k) {
            used[i] = true;
            deg[i] = 0;
            now++;
            for(auto to : g[i]) {
                deg[to]--;
            }
        }
    }
    int remain_edge = 0;
    for(int i = 0; i < n; ++i) {
        remain_edge += deg[i];
    }
    remain_edge /= 2;
    if(now > k) {
        cout << "Impossible" << endl;
        return 0;
    }

    vector<int> ord;
    {
        vector<set<int>> g2(n - now);
        deg.assign(n - now, 0);
        int x = 0;
        vector<int> idx(n);
        vector<pii> deg_idx;
        for(int i = 0; i < n; ++i) {
            if(used[i]) continue;
            deg[x] = 0;
            idx[i] = x++;
        }
        for(int i = 0; i < n; ++i) {
            if(used[i]) continue;
            for(auto to : g[i]) {
                if(!used[to]) {
                    deg[idx[i]]++;
                    g2[idx[i]].insert(idx[to]);
                }
            }
            deg_idx.emplace_back(deg[idx[i]], idx[i]);
        }
        sort(rbegin(deg_idx), rend(deg_idx));
        for(auto& p : deg_idx) {
            ord.push_back(p.second);
        }
        g = move(g2);
        n = g.size();
    }

    int ans = k + 1;
    function<void(int)> solve = [&](int i) {
        if(i == n) {
            ans = min(ans, now);
            return;
        }
        const int r = min(k - now, n - i);
        const int v = ord[i];
        if(now >= ans || now > k || remain_edge > r * deg[v]) {
            return;
        };
        if(g[v].empty()) {
            solve(i + 1);
        } else {
            // use
            if((int)g[v].size() > 1) {
                now++;
                vector<int> erased;
                for(auto to : g[v]) {
                    erased.push_back(to);
                    g[to].erase(v);
                }
                remain_edge -= (int)g[v].size();
                g[v].clear();
                solve(i + 1);
                for(auto to : erased) {
                    g[v].insert(to);
                    g[to].insert(v);
                }
                remain_edge += (int)g[v].size();
                now--;
            }

            // not use
            if(now + (int)g[v].size() <= min(ans, k)) {
                now += g[v].size();
                vector<int> erased;
                vector<set<int>> erased2;
                set<pii> erased_edges;
                for(auto to : g[v]) {
                    erased.push_back(to);
                    set<int> tmp;
                    for(auto to2 : g[to]) {
                        tmp.insert(to2);
                        erased_edges.emplace(min(to, to2), max(to, to2));
                        if(to2 != v) {
                            g[to2].erase(to);
                        }
                    }
                    g[to].clear();
                    erased2.push_back(move(tmp));
                }
                remain_edge -= erased_edges.size();
                g[v].clear();
                solve(i + 1);
                for(int j = 0; j < (int)erased.size(); ++j) {
                    g[v].insert(erased[j]);
                    for(auto to2 : erased2[j]) {
                        g[to2].insert(erased[j]);
                    }
                    g[erased[j]] = move(erased2[j]);
                    remain_edge += g[erased[j]].size();
                }
                now -= g[v].size();
            }
        }
    };
    solve(0);

    if(ans > k) {
        cout << "Impossible" << endl;
    } else {
        cout << ans << endl;
    }
}

感想

実装方針失敗してそう．本番ならもうちょっと高速化サボって実装をスマートにしたほうが良い？
正直枝刈りの実行時間なんて僕には見積もれないんですが……．枝刈り対策どうすればええんや……．毎回お気持ちで書いてしまっていてつらい．

AOJ 2541 Magical Bridges

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2541

解法

特別な辺のコストを定めた時，それぞれの最短経路長がすぐに求まってくれないと困る．
そこで特別な辺は100個しか無いことを利用し，とりあえず以下のダイクストラをやる．

d[i][j] := ゴールから頂点 i まで特別な辺を j 回使ったときの最短経路長

これで特別な辺のコストを定めたときに，各 d[i][j] については O(1) でコストが求まる．

あとはいくつ使うか 100 * 100 通り全部試して全探索する．
ただし，d[s1][i]，d[s2][j] を使うときは，そのパスが最短経路になっていなければならない．
なので，前計算で d[s][i] + x * i = min(d[s][j] + x * j) forall j が成り立つ特別な辺のコスト x の範囲 l[i], r[i] を求めておく．

最後に 100 * 100 通り全部試し，それぞれで一番いいところを二分探索する．
d[s1][i], d[s2][j] を使うなら max(l1[i], l2[j]) から min(r1[i], r2[j]) の範囲になる．
これで範囲を区切ると，各探索区間では求めたいコスト差に単調性があるので二分探索できるということである．

ソースコード

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

constexpr ll inf = 1e18;

struct edge {
    int to;
    ll cost;
    bool magical;
};

using edges = vector<edge>;
using graph = vector<edges>;

int main() {
    int n, m, s1, s2, t;
    while(cin >> n >> m >> s1 >> s2 >> t, n) {
        s1--; s2--; t--;
        graph g(n);
        int mag_num = 0;
        for(int i = 0; i < m; ++i) {
            int a, b;
            string w;
            cin >> a >> b >> w;
            ll c = (w == "x" ? 0 : stoi(w));
            mag_num += (w == "x");
            bool mag = w == "x";
            g[a - 1].push_back(edge{b - 1, c, mag});
            g[b - 1].push_back(edge{a - 1, c, mag});
        }

        vector<vector<ll>> d(n, vector<ll>(mag_num + 1, inf));
        d[t][0] = 0;
        using node = tuple<ll, int, int>;
        priority_queue<node, vector<node>, greater<node>> que;
        que.emplace(0, t, 0);
        while(!que.empty()) {
            ll cur_d;
            int v, mag_cnt;
            tie(cur_d, v, mag_cnt) = que.top();
            que.pop();
            if(cur_d > d[v][mag_cnt]) continue;
            for(auto& e : g[v]) {
                const ll nxt_d = cur_d + e.cost;
                const int nxt_cnt = mag_cnt + e.magical;
                if(nxt_cnt > mag_num || d[e.to][nxt_cnt] <= nxt_d) continue;
                que.emplace(nxt_d, e.to, nxt_cnt);
                d[e.to][nxt_cnt] = nxt_d;
            }
        }

        vector<ll> l1(mag_num + 1), r1(mag_num + 1, inf); // [l1, r1]
        vector<ll> l2(mag_num + 1), r2(mag_num + 1, inf);
        auto calc_lr = [mag_num, &d](int s, vector<ll>& l, vector<ll>& r) {
            for(int i = 0; i < mag_num + 1; ++i) {
                for(int j = 0; j < mag_num + 1; ++j) {
                    if(i == j) continue;
                    if(j - i < 0) {
                        ll t = 0;
                        if(d[s][i] - d[s][j] <= 0) {
                            t = (d[s][i] - d[s][j]) / (j - i);
                        } else {
                            t = -1;
                        }
                        r[i] = min(r[i], t);
                    } else {
                        l[i] = max(l[i], (d[s][i] - d[s][j] + (j - i) - 1) / (j - i));
                    }
                }
            }
        };
        calc_lr(s1, l1, r1);
        calc_lr(s2, l2, r2);

        ll ans = inf;
        for(int i = 0; i < mag_num + 1; ++i) {
            for(int j = 0; j < mag_num + 1; ++j) {
                if(d[s1][i] == inf || d[s2][j] == inf) continue;
                ll lb = max(l1[i], l2[j]);
                ll ub = min(r1[i], r2[j]);
                if(lb > ub) continue;
                while(ub - lb > 1) {
                    const ll mid = (lb + ub) / 2;
                    ll c1 = d[s1][i] + mid * i;
                    ll c2 = d[s2][j] + mid * j;
                    if(c1 < c2) {
                        if(i < j) ub = mid;
                        else      lb = mid;
                    } else {
                        if(i < j) lb = mid;
                        else      ub = mid;
                    }
                }
                ll c1 = abs(d[s1][i] + lb * i - (d[s2][j] + lb * j));
                ll c2 = abs(d[s1][i] + ub * i - (d[s2][j] + ub * j));
                ans = min({ans, c1, c2});
            }
        }
        cout << ans << endl;
    }
}

感想

二分探索パートはなんかサボれる気がしますね．こういう二分探索はバグらせちゃいそうなのでできるだけ避けたい．
候補の点を予め求めてそれぞれについて最短の d[s][i] を求めておく，とか？
でも直線が交差するところは結局変な計算をするのでめんどくさいかなあ．うーん．

AOJ 2434 Audition

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2434

解法

他のそれぞれのアイドルの得点の期待値は，3種類の要素の得点の確率変数をX, Y, Z とすると E[X + Y + Z] だが，これは期待値の線形性から E[X] + E[Y] + E[Z] と等しい．
したがって，それぞれの要素は独立に計算して良い．
すると以下の DP が生える．

dp[i][j][k] := i 人目のアイドルまで見て，担当アイドルが j 回その要素でアピールしたときに，担当アイドルが k 位になる確率

トップ3か最下位しか興味がなく，最下位は簡単に（分けて）求められるので，k = 0, 1, 2, worst の4種類だけ考えれば良い．

それぞれの要素についてこの dp テーブルを求めたら，あとは回数の振り分けを O(m^2) でやって，一番いいやつを選ぶ．この部分は O(m) にもできるけどやる必要はない．

ソースコード

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n, m;
    cin >> n >> m;
    vector<int> vi(n), da(n), vo(n);
    for(int i = 0; i < n; ++i) {
        cin >> vi[i] >> da[i] >> vo[i];
    }

    vector<double> p(m + 1, 0);
    p[0] = 1;
    for(int i = 0; i < m; ++i) {
        vector<double> nxt(m + 1);
        for(int i = 0; i <= m; ++i) {
            nxt[i] += p[i] * 2 / 3.0;
            if(i != m) {
                nxt[i + 1] += p[i] / 3.0;
            }
        }
        p = move(nxt);
    }
    for(int i = 0; i < m; ++i) {
        p[i + 1] += p[i];
    }

    auto calc = [&](vector<int> const& v) {
        vector<vector<double>> dp(m + 1, vector<double>(4));
        for(int i = 0; i < m + 1; ++i) {
            dp[i][0] = dp[i][3] = 1;
        }
        for(int i = 1; i < n; ++i) {
            vector<vector<double>> nxt(m + 1, vector<double>(4));
            for(int j = 0; j <= m; ++j) {
                const int t = v[0] * j;
                const int cnt = (t + v[i] - 1) / v[i];
                const double win = p[min(m, cnt - 1)];
                for(int k = 0; k < 3; ++k) {
                    nxt[j][k] += (cnt == 0 ? 0 : dp[j][k] * win); // win
                    if(k + 1 < 3) {
                        nxt[j][k + 1] += (cnt == 0 ? dp[j][k] : dp[j][k] * (1 - win)); // lose
                    }
                }
                nxt[j][3] += dp[j][3] * (1 - win);
            }
            dp = move(nxt);
        }
        return dp;
    };
    auto dp1 = calc(vi);
    auto dp2 = calc(da);
    auto dp3 = calc(vo);

    double ans = -1e9;
    for(int i = 0; i <= m; ++i) {
        for(int j = 0; j + i <= m; ++j) {
            int k = m - i - j;
            double t = (dp1[i][0] + dp1[i][1] + dp1[i][2]) * 5;
            t += (dp2[j][0] + dp2[j][1] + dp2[j][2]) * 3;
            t += (dp3[k][0] + dp3[k][1] + dp3[k][2]) * 2;
            t -= (dp1[i][3] + dp2[j][3] + dp3[k][3]);
            ans = max(ans, t);
        }
    }

    cout << fixed << setprecision(10) << ans << endl;
}

AOJ 1324 Round Trip

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=1324

解法

最短経路（それはそう）．
同じ高さの街が高々10しかないので，どの街に行ったか 2 ^ 10 で表現したい．
どの街に行ったかを保存しつつ最短経路をやりたいね，となる．
それにはまず，元のグラフの辺の向きをすべて逆に張った逆グラフをもう一つ用意する．
すると，戻りの条件が行きと全くおなじになる（低い方から高い方へ）．
ここまでやって以下の最短路を解く．

d[i][j][S] := 行きのパスの最後が i ，戻りのパスの最後（厳密に言うと先頭）が j であり，max(h[i], h[j]) の街の集合の使用状態がSであるときの最小コスト

ただし h[i] は街 i の高さ．
遷移するときは，max(h[i], h[j]) 以上のところしか行けないようにする．
遷移する条件は i を動かすときも j を動かすときも同じ．

計算量は O(n ^ 2 * 2^10 * log(...)) 程度．複数テストケースだから若干怖いけど，通るらしい．

ソースコード

#include <bits/stdc++.h>
using namespace std;

constexpr int inf = 1e9;

struct edge {
    int to, cost;
};

using edges = vector<edge>;
using graph = vector<edges>;

int main() {
    int n, m;
    while(cin >> n >> m, n) {
        vector<int> height(n), fee(n), bit(n);
        height[n - 1] = 1000;
        vector<int> h_cnt(1000);
        for(int i = 1; i < n - 1; ++i) {
            cin >> fee[i] >> height[i];
            bit[i] = h_cnt[height[i]]++;
        }
        graph g(n), rg(n);
        for(int i = 0; i < m; ++i) {
            int a, b, c;
            cin >> a >> b >> c;
            g[a - 1].push_back(edge{b - 1, c});
            rg[b - 1].push_back(edge{a - 1, c});
        }

        vector<vector<vector<int>>> d(n, vector<vector<int>>(n, vector<int>(1 << 10, inf)));
        d[0][0][1] = 0;
        using state = tuple<int, int, int, int>;
        priority_queue<state, vector<state>, greater<state>> que;
        que.emplace(0, 0, 0, 1);
        while(!que.empty()) {
            int cur_d, u, v, S;
            tie(cur_d, u, v, S) = que.top();
            que.pop();
            const int cur_h = max(height[u], height[v]);
            for(auto& e : g[u]) {
                if(height[e.to] < cur_h) continue;
                const int nS = (1 << bit[e.to]) | (height[e.to] > cur_h ? 0 : S);
                const int nxt_d = cur_d + (height[e.to] == cur_h && (S & (1 << bit[e.to])) ? 0 : fee[e.to]) + e.cost;
                if(nxt_d < d[e.to][v][nS]) {
                    d[e.to][v][nS] = nxt_d;
                    que.emplace(nxt_d, e.to, v, nS);
                }
            }
            for(auto& e : rg[v]) {
                if(height[e.to] < cur_h) continue;
                const int nS = (1 << bit[e.to]) | (height[e.to] > cur_h ? 0 : S);
                const int nxt_d = cur_d + (height[e.to] == cur_h && (S & (1 << bit[e.to])) ? 0 : fee[e.to]) + e.cost;
                if(d[u][e.to][nS] > nxt_d) {
                    d[u][e.to][nS] = nxt_d;
                    que.emplace(nxt_d, u, e.to, nS);
                }
            }
        }

        cout << (d[n - 1][n - 1][1] == inf ? -1 : d[n - 1][n - 1][1]) << endl;
    }
}

感想

実装重い系かと思いきや，そんなことはなかった．

Codeforces Round #176 (Div.1) E. Ladies' Shop

問題文

http://codeforces.com/problemset/problem/286/E

解法

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 すると通る．