【pythonでABC137を解説】D - Summer Vacation

問題概要

問題ページ

問題文

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

あなたは、これらの中から \(1\) 日に \(1\) 件まで選んで請け、働くことができます。

ただし、請けたことのある日雇いアルバイトは選べません。

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

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

制約

• 入力は全て整数である。
• \(1 \leq N \leq 10^5\)
• \(1 \leq M \leq 10^5\)
• \(1 \leq A_i \leq 10^5\)
• \(1 \leq B_i \leq 10^4\)

問題の考察

優先度付きキュー（pythonだとheapq）の使い方が分かれば解ける問題です。

状況を整理するために「入力例2」を考えてみましょう。

この問題のポイントは「制限の厳しいものから埋めていく」です。

図を見れば、この例の場合には\(2\)日後の選択するバイトが一番制限がきついことが分かります。

\(2\)日後に選択できるアルバイトは\(1\)日後でも今日でも選択できますが、今日選択できるアルバイトは\(2\)日後に選択できるとは限りません。

\(2\)日後に選択できて報酬が一番高いバイト、\(1\)日後に選択できて報酬が一番高いバイト、今日選択できて報酬が一番高いバイト・・・・という順番で埋めていけば良いことが分かります。

たびすけ

制限が厳しいところから決定していくと簡単に解けるよ！

次に、どのようにコーディングするかですが、次のようになります。

ここでリスト\(t\)から一番高い報酬を選択する時に「heapq」を使います。

通常であればリスト内の最小値を取得する時に\(O(N)\)の計算量がかかりますが、「heapq」は\(O(\log N)\)で取得できるライブラリです。

注意点としてpythonのheapqは「最小値しか取得できない」ので、最大値を取得したい場合には追加の際に\(b\)なら\(-b\)のように符号を逆にしてから追加してあげる必要があります。

heapqを使えば全体の計算量が\(O(\log N \times M)\)となり制限時間内に解答できました。

たびすけ

「heapq」は競技プログラミングで頻出なので忘れていたら復習して覚えましょう！

ACコード

import sys
import heapq
    
    
def solve():
    input = sys.stdin.readline
    mod = 10 ** 9 + 7
    n, m = list(map(int, input().rstrip('\n').split()))
    ab = [list(map(int, input().rstrip('\n').split())) for _ in range(n)]
    heapq.heapify(ab)
    total = 0
    t = []
    heapq.heapify(t)
    for i in range(m):
        while len(ab):
            a, b = heapq.heappop(ab)
            if i + 1 >= a:
                heapq.heappush(t, [-b, a])
            else:
                heapq.heappush(ab, [a, b])
                break
        if len(t):
            total -= heapq.heappop(t)[0]
    print(total)
    
    
if __name__ == '__main__':
    solve()