hogecoder

tsutaj 競技プログラミングの記録

Codeforces Round #378 C: Epidemic in Monstropolis

問題概要

原文 → Problem - 733C - Codeforces

サイズ nの配列 aと、サイズ kの配列 bが与えられる。配列 aに対して、隣り合う要素が自分より小さい場合、自分にその要素の値を加算し、その要素を削除する操作(つまりマージ)が可能である。この操作を繰り返して配列 aを配列 bと一致させることは可能か。もし可能な場合、操作の過程(何番目の要素を左or右とマージさせたか)も出力すること。

解説

この問題はコーナーケースが多く、気をつけないとWAやTLE祭りになってしまいますね。私もwhileループ抜けずにTLE問題に直面しました。

まず、 \sum_{i=1}^{n} a_{i} = \sum_{i=1}^{k} b_{i}を満たしていなければ、明らかに達成できません。この判定は最初にすべきでしょう。

これを満たしている場合、配列 aを適当な k個の区間に分けることを考えます。具体的に言うと、 i個目の区間の要素の総和が b_iと等しくなるように分けます。(このような分け方が存在しなければ達成不可能です)

区間内でどことどこをマージするかですが、区間内で値が最大の要素で隣とマージ可能な要素を選んでマージするのが最適で、なぜならそうすることでその要素が常に区間内max(しかも単独トップ)を保ち続けるからです。そのような要素がなければ残念ながら達成不可能なため、 \verb$"NO"$を出力する必要があります。ちなみに、区間内の要素が1になるまでマージし続けてから次の区間に行ったほうが簡潔に書けます。

達成不可能な条件が多いのでまとめると、

  • 配列 aの総和と配列 bの総和が異なるとき
  • 配列 aを適当な区間に分け、 i個目の区間総和を b_iと等しくする方法がないとき
  • 区間内最大の要素で、隣とマージ可能な要素がないとき

となります。1つでも当てはまれば達成不可能です。

ソースコード

signed main() {
    int sum_a = 0, sum_b = 0;
    int n, k; cin >> n; vector<int> a(n);
    rep(i,0,n) {cin >> a[i]; sum_a += a[i];}
    cin >> k; vector<int> b(k);
    rep(i,0,k) {cin >> b[i]; sum_b += b[i];}
    if(sum_a != sum_b) {
        cout << "NO" << endl;
        return 0;
    }

    vector< pair<int, char> > ans;
    int x = 0;
    rep(i,0,k) {
        int temp = 0;
        vector<int> v;
        while(x != n && temp < b[i]) {
            v.pb(a[x]);
            temp += a[x++];
        }
        if(temp != b[i]) {
            cout << "NO" << endl;
            return 0;
        }

        while(v.size() != 1) {
            bool ok = false;
            int ma = *max_element(v.begin(), v.end());
            rep(j,0,v.size()) {
                if(v[j] == ma) {
                    if(j != 0 && v[j-1] < v[j]) {
                        v[j-1] += v[j];
                        ans.pb(make_pair(j+i, 'L'));
                        v.erase(v.begin() + j);
                        ok = true; break;
                    }
                    else if(j != v.size()-1 && v[j+1] < v[j]) {
                        v[j+1] += v[j];
                        ans.pb(make_pair(j+i, 'R'));
                        v.erase(v.begin() + j);
                        ok = true; break;
                    }
                }
            }
            if(!ok) {
                cout << "NO" << endl;
                return 0;
            }
        }
    }
    cout << "YES" << endl;
    rep(i,0,ans.size())
        printf("%lld %c\n", ans[i].fr+1, ans[i].sc);
    return 0;
}

ちなみにこれはコンテスト終了後に書きなおしたもので、終了直後にACしたものはもっと汚いです(下のリンク参照、こりゃあ間違い誘発するわって感じ・・・)

バグを誘発させる事のないように書くことって大事ですね。

Submission #21956334 - Codeforces