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

hamayanhamayan's blog

Pair of Balls [AtCoder Beginner Contest 216 D]

https://atcoder.jp/contests/abc216/tasks/abc216_d

解説

https://atcoder.jp/contests/abc216/submissions/25449941

貪欲法というかシミュレーションで解く問題。
今回は目標を達成するには、ボールが取り出せるならどんな順番でもいいので取り出していって、
最後に全部取り出せれば達成可能と判定する方針で実装していく。
なので、高速にこのことをシミュレーションしていくことを考える。

このように選択肢があるように見えて、どんな方針でやっても最適解になるような問題がある。
これも正当性判断が難しいですが、とりあえず選択肢が2つ以上あるときに、操作の順番を変えることで
有利な状況が作れそうかというのを例を作りながら確認していくしかないと思う。
点数を見て、実装が大変な系かなという判断を下したのかも。

高速シミュレーション

ボールを取り出す際に毎回筒を確認して、同じボールがあるかを探していく方針はN回取り出す必要があって、
雑に計算すると毎回M個の筒を確認する必要があるので、O(NM)は間に合わない感じがする。
なので、上手くメモを残しながら取り出せるボールというのを高速に取ってくるようにする。
以下、自分の実装について説明する。

以下のようなデータ構造を用意する。

que := 「現在取り出せる同色のボールの組」を集めたキュー
available[i] := 現在色iのボールが先頭にある筒の番号

例えばi番目の筒の先頭がA[i].front()であるとすると、
available[ A[i].front() ] = i
という風に更新することができる。
だが、更新時に既に値があった場合は、他にも筒の先頭に同じ色のボールがあるということになるので、
取り出すためにqueというキューに筒の番号を追加する
que.push({ i, available[ A[i].front() ] })
自分の番号iと、前もってメモっておいた同色が先頭にある筒の番号available[ A[i].front() ]をキューに入れている感じ。
この実装を

setAvailable(ai) := 筒A[ai]の先頭を評価する

として実装している。
キューを1つずつ処理していくことで同色の2組を削除していくが、2つの筒からボールが取り出されるので
その次のボールをsetAvailableでavailable配列の更新やqueへの追加を行っていく。
これでボールを取り出して、差分として新たに2つのボールが出てくるので、それと先頭が一致しているものがないか確認というのを
上手いこと差分だけを計算することで高速に処理していく。

最後に

シミュレーション後にすべての筒が空になっていればYes。
筒の先頭を評価するときに空の場合は無視することに注意。
あと、筒の管理にdequeを使っている。

int N, M;
deque<int> A[201010];
int available[201010];
//---------------------------------------------------------------------------------------------------
queue<pair<int, int>> que;
void setAvailable(int ai) {
    if (A[ai].empty()) return;
    if (0 <= available[A[ai].front()]) {
        que.push({ available[A[ai].front()] , ai });
    }
    else available[A[ai].front()] = ai;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> M;
    rep(i, 0, M) {
        int k; cin >> k;
        rep(j, 0, k) {
            int a; cin >> a;
            a--;
            A[i].push_back(a);
        }
    }
    rep(i, 0, N) available[i] = -1;

    
    rep(i, 0, M) setAvailable(i); {
        
    }
    while (!que.empty()) {
        int i1, i2;
        tie(i1, i2) = que.front();
        que.pop();

        A[i1].pop_front();
        setAvailable(i1);
        A[i2].pop_front();
        setAvailable(i2);
    }

    bool ok = true;
    rep(i, 0, M) if (!A[i].empty()) ok = false;

    if (ok) cout << "Yes" << endl;
    else cout << "No" << endl;
}