priority_queueの入門書<競プロ編> #C++

はじめに

どうも、皆さんこんにちは、普段C++とイチャイチャしてるJinです(Twitter: @CPP_IS_GOD)。先日行われたABC141でpriority_queueを使う問題が出題されました。個人的にpriority_queueはいつか学ばないとなぁ～とは思っていたので、今回一回整理してみようと思います。

対象読者

priority_queueの名前は知ってた人
queueは使ったことあるけどpriority_queueは使ったことない人

queue(キュー)とは

では、まずqueueについて説明しましょう。queueはコンテナの一種で、ほとんどの言語では標準サポートされています。queueはよく「レストランで待っている行列」と例えられます。理由は簡単です。queueコンテナは「レストランで待っている行列」のように振る舞うからです。英語が公用語の国では行列の近くに"queue"と書かれていることがよくあります。
本題に戻ると、queueには大きく分けて2つの操作をすることができます。

追加する
取り出す

レストランの行列もこの2つの操作で説明することができます。「追加する」は"予約する"に対応し、「取り出す」は"ウェイターに呼ばれて店内に入る"に対応します。
ね、わかりやすいでしょ？(脅し)
ちなみに、先に追加した要素が先に取り出されることからqueueはFIFO(First-In, First-Out)と呼ばれます。逆に後に追加した要素から取り出されるコンテナもあり(stackという)、それはLIFO(Last-In, First-Out)と言われています。
説明だけだと分かりづらいので実際にC++で書いてみましょう。

queue<int> q;//キューコンテナの生成
    q.push(1);//1を追加
    q.push(2);//2を追加
    q.push(3);//3を追加

    cout << q.front() << endl;//1を最初に追加したので当然1が出力される
    q.pop();
    /*
    C++の場合、front()は次の要素の値を返すことしかせず、取り出すにはpop()を別途に呼び出す必要がある。
    他の言語の場合、要素の取得と取り出しを1つのメソッドで処理してくれることが多い。
    */
    cout << q.front() << endl;//もちろん2が出力される
    q.pop();
    cout << q.front() << endl;//もちろん3が出力される
    q.pop();

これで大体のqueueの挙動が理解できたと思います。

priority_queue(優先度付きキュー)とは

さて、ここからが本題です。queueを理解した皆さんならこう思うはずでしょう、「なんかしらの優先度をつけてその順番で取り出したい」と(だいぶ無理やり)！それを実現したのがpriority_queueなのです。
エコノミークラスの人が先に空港の搭乗口に並んでいたのにファーストクラスから先に機内に入れる(ファーストクラスのほうがエコノミークラスより優先度が上なので)経済力の差を見せつけられるシステムに非常によく似ていますね！
そんな話はおいといて、デフォルトだとpriority_queueは値を昇順にソートされて取り出されます。ここで例を示しましょう。

priority_queue<int> q;//生成方法はqueueと同じ
    q.push(1);
    q.push(3);
    q.push(2);

    cout << q.top() << endl;
    /*
    queueだったら1が出力されるはずだが、昇順にソートされてから取り出されるので3が出力される！！
    3「🤣✌」
    1「解せぬ...」
    (次の要素へのアクセスはfront()ではなくtop()なので注意)
    */
    q.pop();
    cout << q.top() << endl;//2が出力
    q.pop();
    cout << q.top() << endl;//1が出力
    q.pop();

各要素内の格差社会を紹介したところで、頭の良い人は次のことを思うでしょう、「要素lと要素r間の優先順位をラムダ式等の関数オブジェクトで定義したらユーザー定義ができるのでは？？」と。
はい、できます。
では、入ってきた値の絶対値で優先順位をつけてみましょう。

auto comp = [](int l, int r)->bool {return (abs(l) < abs(r)); };//比較用のラムダ式を生成
    //[注意]もし、lの優先度が低い場合はtrueを返すようにする。

    priority_queue<int, std::vector<int>, decltype(comp) > q(comp);
    /*
    優先度をユーザー定義にする場合、テンプレート引数に内部で使うコンテナクラスと関数オブジェクトの型を渡してあげる必要がある。
    */

    q.push(5);
    q.push(-6);
    q.push(1);

    cout << q.top() << endl;//-6の絶対値は6であり5より大きいので-6が出力される！！
    q.pop();
    cout << q.top() << endl;//5が出力
    q.pop();
    cout << q.top() << endl;//1が出力
    q.pop();

これでだいたいのpriority_queueの使い方は理解できたと思います。

実際に競プロの問題で使ってみよう！

というわけで、実際の競プロの問題でpriority_queueを使ってみましょう。例題としてABC141-Dを使わせていただきます。

問題文

高橋くんはN個の品物を1個ずつ順番に買う予定です。
i番目に買う品物の値段は A_i円です。
高橋くんはM枚の割引券を持っています。
品物を買う際に割引券を好きな枚数使うことができます。
X円の品物を買う際に Y枚の割引券を使った場合、その品物を X/2^Y円(小数点以下切り捨て)で買うことができます。
最小で何円あれば全ての品物を買うことができるでしょうか。

(制約等の細かい問題文はリンクを飛んで確認してください)

方針

⌊X/2^Y⌋ = ⌊⌊⌊X/2⌋/2⌋/2⌋...
と式変形させることができます。(詳しい説明は解説放送を参照してください)
よって各値段の最大値をM回、/=2をすればよい(数が大きいほうが、1/2になるときに浮く金の値が増えるため)。
これを貪欲法でやろうとすると計算量が$O(NM)$になるのでTLE(時間制限超え)になってしまいます。
だったら、priority_queueで昇順ソートを内部でしながらpush&popをすればいいのでは？というのが方針です。
なぜなら、priority_queueの"値の追加"と"値の削除"は共に計算量が$O(log (size))$になるのでTLEを防げるからです。

実装

ここまで来たら、賢い皆さんなら実装できるのでソースコードだけ置きます。

int N, M; cin >> N >> M;//入力の受け取り
    priority_queue<int> q;//キューの作成(デフォルトで昇順ソートしてくれるので特にいじらなくていいです)
    rep(i, N) {//各要素をキューに追加
        int A; cin >> A;
        q.push(A);
    }
    rep(i, M) {//最大値を取り出して1/2にしてまた追加(M個割引券があるのでM回回す)
        int A = q.top(); q.pop();
        q.push(A / 2);
    }
    ll ans = 0;//念の為long longにしましょう
    while (!q.empty()) {//割引券を前のループで使い切ったのでキューが空になるまでキューの値をansに追加
        ans += q.top(); q.pop();
    }

    cout << ans;//最後に答えを出力

僕のACコード: https://atcoder.jp/contests/abc141/submissions/7587149

参考文献

https://www.youtube.com/watch?v=fHZhDUzhzN0 (解説生放送)
https://qiita.com/y_shindoh/items/17d9fa334a2cb8e74bfa (priority_queueについての解説記事)

最後に

皆様、楽しい競プロライフを！
ツイッターのフォローお願いします(@CPP_IS_GOD)！！(最後にバッキバキの宣伝して終わるスタイル)