けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 137 D - Summer Vacation (水色, 400 点)

AtCoder AtCoder400点 ABC-D 水色diff Greedy priority_queue データ構造 Greedy:端から順に決まっていく 後ろから解く 探索順序を工夫して解く 最適化テク:解を変形していく(最適性を失わずに) マトロイド 勤怠を題材とした問題

これは難しい！！！　誘惑されそうな嘘解法がたくさんある！！

問題概要

N 件の日雇いアルバイトがあります。 i 件目の日雇いアルバイトを請けて働くと、その A_{i} 日後に報酬 B_i が得られます。

あなたは、これらの中から 1 日に 1 件まで選んで請け、働くことができます。ただし、請けたことのある日雇いアルバイトは選べません。

今日から M 日後まで ( M 日後を含む) に得られる報酬の合計の最大値を求めてください。

なお、日雇いアルバイトは今日から請けて働くことができます。

制約

  • 1 \le N, M \le 10^{5}
  • 1 \le A_i \le 10^{5}
  • 1 \le B_i \le 10^{4}

考えたこと

一見すると自然に思えてしまう嘘解法がたくさんある！！

たとえば、「報酬が多いアルバイトから順にやる」という Greedy に対しては、 M = 2, (A, B) = (1, 5), (2, 1) の場合が反例となる。

ほかにも、「報酬を受け取れるまでの期間が長いアルバイトから順にやる (先に終わらせておきたい)」という Greedy に対しても、 M = 2, (A, B) = (1, 5), (1, 5), (2, 1) の場合が反例となる。

うまい探索順序を見出す

手の付け所が難しい問題だと思う。このような問題では、探索順序を適切に定めることが重要かなと。

さて、この問題では、最終日の方から考えていくとうまくいく。まず、締め切り前日である M-1 日目に選べるアルバイトはそもそも「 A_{i} = 1 であるもの」しかないのだ。こういう「ここにしか入れられない」というものから順に考えていくのは有効だ。というわけで、次のことが言える。

「まず M-1 日目には、 A_{i} = 1 であるアルバイトのうち、 B_{i} が最も大きいアルバイト ( p とする) を入れてしまってよい」

このように言えることの証明は、いつも通りの議論でできる。つまり、アルバイト p なしの解が得られているときに、その解を悪化させることなく p を挿入することができる or 入れ替えることができるのだ。

M-1, M-2, \dots, 0 日目と遡る

次に M-2 日目を考える。まったく同様にして、

  • 残っているアルバイトのうち
  • A_{i} \le 2 であるアルバイトであって
  • B_{i} が最大のもの

を入れてしまってよいということがいえる。以降これを繰り返すと、次の解法が導かれる。

i = M-1, M-2, \dots, 0 に対して

  • 残っているアルバイトのうち
  • A_{i} \le M - i であるアルバイトであって
  • B_{i} が最大のもの

を i 日目に入れていく。ただし存在しない場合にはスキップする。

データ構造を用いて高速化

以上の解法は、単純に実装すると O(N^{2}) の計算量となるので高速化が必要だ。ここではヒープとよばれるデータ構造を用いて高速化する。ヒープは次のことができる。

  • ヒープに要素 v を挿入する ( O(\log N))
  • ヒープに含まれる要素のうち、最大値を取得する ( O(1))
  • ヒープに含まれる要素のうち、最大のものを削除する ( O(\log N)

ヒープは、C++ では priority_queue、Python3 では heapq が使える。ヒープを活用することで、計算量は O((M + N) \log N) となる。　

コード

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

int main() {
    int N, M;
    cin >> N >> M;

    // AtoB[i] := A が i であるような B の集まり
    vector<vector<int>> AtoB(M+1);
    for (int i = 0; i < N; ++i) {
        int A, B;
        cin >> A >> B;
        if (A > M) continue;
        AtoB[A].push_back(B);
    }

    long long result = 0;
    priority_queue<long long> que;
    for (int i = M-1; i >= 0; --i) {
        for (auto B : AtoB[M-i]) que.push(B);
        if (que.empty()) continue;
        result += que.top(); // キューの最大値を足す
        que.pop(); // キューの最大要素を削除
    }
    cout << result << endl;
}