はまやんはまやんはまやん

hamayanhamayan's blog

Happy Birthday! 2 [京セラプログラミングコンテスト2021(AtCoder Beginner Contest 200) D]

https://atcoder.jp/contests/abc200/tasks/abc200_d

解説

https://atcoder.jp/contests/abc200/submissions/22441528

DPして経路復元する。
(公式解説は鳩ノ巣原理です…なるほど…)

DP

DPについては組合せDPが分かっていれば中身はそれほど難しくない。

dp[i][mo][isSelected] := i番目の数まで利用していて総和を200で割った余りがmoで、1つでも数が選択されているかどうかがisSelectedである状態を満たす組み合わせ

この組み合わせを計算してdp[N][0][1]~dp[N][199][1]のいずれかが2以上になっていればそのmodを満たす選択が2通りあるということなので、
その選択を抜き出してくれば答えになる。
組合せ数が膨大になるかもしれないので、更新時に最大値が2になるようにトリムしておく。
今回は経路復元も必要になるので、経路復元用の配列も用意しておこう。

pre[i][mo][isSelected] := dp[i][mo][isSelected]の状態からさかのぼる状態を保持する

形式は{ mo', isSelected' * 2 + i番目を選択したか }としていて、dp[i][mo][isSelected]からdp[i-1][mo'][isSelected']へ戻す。
i番目を選択したかは経路復元時に選択したかどうかが分からないと答えとして採用するかが分からなかったので入れている。

経路復元

経路復元がやや厄介である。
pre配列を作るときにさかのぼる状態を最大2つ保持するようにしておく。
そして、B,Cを作る場合に、このpre配列の状態を

  • Bなら戻る状態のうち先にあるもの
  • Cなら戻る状態のうち後にあるもの

を必ず選択するようにしていけば、違う経路が2つ得られる。
あとは、それぞれ持ってきて答える。

int N, A[201];
//---------------------------------------------------------------------------------------------------
int dp[201][201][2];
vector<pair<int, int>> pre[201][201][2];
void doDP() {
    dp[0][0][0] = 1;
    pre[0][0][0].push_back({ -1, -1 });
    rep(i, 0, N) rep(mo, 0, 200) rep(isSelected, 0, 2) if (0 < pre[i][mo][isSelected].size()) {
        dp[i + 1][mo][isSelected] += dp[i][mo][isSelected];
        chmin(dp[i + 1][mo][isSelected], 2);

        if (pre[i + 1][mo][isSelected].size() < 2) {
            pre[i + 1][mo][isSelected].push_back({ mo, isSelected * 2 + 0 });
        }

        int mo2 = (mo + A[i]) % 200;

        dp[i + 1][mo2][1] += dp[i][mo][isSelected];
        chmin(dp[i + 1][mo2][1], 2);
        if (pre[i + 1][mo2][1].size() < 2) {
            pre[i + 1][mo2][1].push_back({ mo, isSelected * 2 + 1 });
        }
    }
}
//---------------------------------------------------------------------------------------------------
vector<int> getAns(int ok_mo, bool isB) {
    vector<int> ans;
    int mo = ok_mo, isSelected = 1;
    rrep(i, N, 1) {
        auto v = pre[i][mo][isSelected];

        int idx = isB ? 0 : v.size() - 1;

        mo = v[idx].first;
        isSelected = v[idx].second / 2;
        int used = v[idx].second % 2;

        if (used) ans.push_back(i);
    }
    reverse(all(ans));
    return ans;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> A[i];

    doDP();
    
    int ok_mo = -1;
    rep(mo, 0, 200) if (dp[N][mo][1] == 2) ok_mo = mo;

    if (ok_mo < 0) {
        cout << "No" << endl;
        return;
    }

    printf("Yes\n");

    auto B = getAns(ok_mo, true);
    auto C = getAns(ok_mo, false);

    int x = B.size();
    printf("%d", x);
    rep(i, 0, x) printf(" %d", B[i]);
    printf("\n");

    int y = C.size();
    printf("%d", y);
    rep(i, 0, y) printf(" %d", C[i]);
    printf("\n");
}