AT274の技術だったりなかったり

あっとのTECH LOG

競プロのこととか技術のこととか。たまに日常を書きます。

ABC137 D - Summer Vacation

問題原文

atcoder.jp

問題要旨

 N個 の仕事があり、 仕事  i A_i 日後に  B_i の報酬をもらえる。 1日にこなせる仕事は1件である。
 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}

解法

ソートして貪欲したくなるが。。。?

ギリギリ報酬がもらえるような仕事をやる?

これは反例が簡単に見つかって、

  •  M = 3
  •  A_1 = 3,  B_1 = 1
  •  A_2 = 2,  B_2 = 100
  •  A_3 = 2,  B_3 = 100

のようなパターンで簡単に落ちる。
ので×。

報酬がもらえるような仕事のうち、最も価値が高いものをやる?

これならさっきの、

  •  M = 3
  •  A_1 = 3,  B_1 = 1
  •  A_2 = 2,  B_2 = 100
  •  A_3 = 2,  B_3 = 100

でも通る。これなら大丈夫そう。
でも、「報酬がもらえるような仕事を、 1日目、2日目、3日目…  M - 1 日目」 と計算していくのはなんかたいへん。

1日目、2日目… ではなく、  M - 1 日目、  M - 2 日目…と考える。

「可能なものからどんどん削っていく」よりは、「可能なものを追加していく」ほうが遥かに簡単。
よって、  Ai = 1, 2, 3 ... と小さい方から範囲を広げていくことを考える。

これで、「報酬がもらえるような仕事」が高速に求められるので、あとはいかに高速に最も価値が高いものを取り出すか。
これは優先度付きキューを使えばいいことが、少し考えるとわかる。
計算量は、 候補の追加  O(N) 回、 候補の中から最大のものを取り出すのを  O(M) 回、 + 優先度付きキューを使うので、  O(M + NlogN) となる。

実装

import heapq
N, M = map(int, input().split())
works = sorted([list(map(int, input().split())) for i in range(N)])

ans = 0
hq = []
i = 0
for limit in range(1, M + 1):
    while (i < N) and (works[i][0] <= limit):
        cost, value = works[i]
        heapq.heappush(hq, -value)
        i += 1

    if hq:
        ans += -heapq.heappop(hq)

print(ans)

感想

綺麗な問題。
かなり好き。

「全体から削っていく」 → 「空に追加していく」の転換は汎用性が高そうな考え方ですね。