F - Erase and Rotate Editorial by MMNMM

公式解説の補足

公式解説におけるシミュレーションを \(O(N)\) で実現する方法について書きます。

次のような問題を解くことになります。

  • 各要素が \(\mathbb{N}\times\{0,1\}\) の元である列 \((p_i,c_i)_{i=1,\ldots,N}\) が与えられます。\(i\) 番目の要素を削除するコストは \(c_i\) です。コストの合計が \(K\) 以下であるような削除で残る列のうち、辞書順で最小であるようなものを求めてください。

これは、stackを用いて貪欲を行うことで解くことが可能です。つまり、次のようなアルゴリズムが答えを与えます。

  • \(i=1,\ldots,N\) と見ていく。 stack の先頭が \((p_i,c_i)\) より大きく、stack の先頭を削除しても合計コストが \(K\) を上回らないなら、 stack の先頭を削除する。これを繰り返して条件を満たさなくなったとき(つまり、stack が空になるか、先頭が \((p_i,c_i)\) より小さいか、合計コストが \(K\) を上回ったとき)、stack に \((p_i,c_i)\) を追加する。

stack から削除される要素は stack に追加された要素なので、これは \(O(N)\) 時間で動作します。

実装は以下のようになります。実際の提出

#include <bits/extc++.h>

int main(){
    using namespace std;
    unsigned long N, K;
    cin >> N >> K;

    const auto simulation{[](auto&& p, unsigned long K){
        vector<pair<unsigned long, unsigned long>> sim_stack;
        p.emplace_back();
        
        unsigned long remain_op = K;
        
        for(const auto& i : p){
            while(!empty(sim_stack) && remain_op >= sim_stack.back().second && sim_stack.back() > i){
                remain_op -= sim_stack.back().second;
                sim_stack.pop_back();
            }
            sim_stack.emplace_back(i);
        }
        sim_stack.pop_back();
        p.pop_back();
        
        vector<unsigned long> ret(size(sim_stack));
        transform(begin(sim_stack), end(sim_stack), begin(ret), [](auto&& i){return i.first;});
        
        return ret;
    }};

    vector<pair<unsigned long, unsigned long>> p(N, {1, 1});
    for(auto&& [i, _] : p)cin >> i;

    auto p_rotated = p;
    const unsigned long front_idx = min_element(begin(p) + N - K, end(p)) - begin(p);
    for(auto i = front_idx; i < N; ++i)p_rotated[i].second = 0;
    rotate(begin(p_rotated), begin(p_rotated) + front_idx, end(p_rotated));
    
    const auto ans = min(simulation(p, K), simulation(p_rotated, K - N + front_idx));
    for(const auto i : ans)cout << i << " ";
    cout << endl;

    return 0;
}

posted:
last update: