PythonのheapqでAtCoderの問題を解くときは、先にヒープへ入れるキーと候補を更新する時系列を決めます。答えが一つの最大値なら値の負数を入れ、全要素へ同じ値を足すなら共通のoffsetへ分離します。仕事の締切は後ろの枠から、開始時刻と退出時刻を持つ商品は前から候補を更新すると、毎回全要素を走査せずに済みます。
同じ優先度付きキューでも、キーの向き、全体更新の表現、候補が有効になる時点、期限切れの境界が変わることを、公式問題の入力とサンプルで追える順番にしています。読了後は、問題文を見て次の4点をコードの前に決められる状態を目指します。
- 一回の操作で最小値と最大値のどちらが必要か。
- 全要素への同じ加算を
base + offsetで表せるか。 - 候補をヒープへ追加する日や時刻はいつか。
- 境界を含む候補を、どの条件で期限切れとして捨てるか。
ABC141 D：最大値を負のキーで毎回更新する
ABC141 Dの公式問題文では、N個の商品にM枚の割引券を使います。価格Xの商品へY枚使った価格は floor(X / 2^Y) です。制約は1 <= N, M <= 10^5、1 <= A_i <= 10^9です。公式EditorialのD – Powerful Discount Ticketsが示すように、同じ商品への複数枚の割引は、一枚ずつ現在価格を半分にする操作として扱えます。
なぜ最大値を選ぶのか
一枚の割引券を価格xへ使うと、価格はfloor(x / 2)になります。価格が大きい商品ほど一枚で減る金額が大きいので、現在の最大価格へ使うのが得です。二つの価格x >= yを比べると、xを半分にする場合の減少量は、yを半分にする場合の減少量以上です。別の商品へ使った解があっても、最後の一枚を最大価格へ交換して合計価格を悪化させないため、各回で最大値を選べます。
公式サンプル1の入力は、商品価格が2, 13, 8、割引券が3枚です。最大価格を選ぶと状態は次のように変わります。
| 割引券を使う前 | 選ぶ価格 | 半分にした後 | 合計 |
|---|---|---|---|
| 2, 13, 8 | 13 | 2, 6, 8 | 16 |
| 2, 6, 8 | 8 | 2, 6, 4 | 12 |
| 2, 6, 4 | 6 | 2, 3, 4 | 9 |
heapqは最小値を取り出すため、価格そのものではなく負の価格を入れます。たとえば価格13は-13となり、負数の中で最小のキーが現在価格の最大値を表します。
import heapq
N, M = map(int, input().split())
heap = [-value for value in map(int, input().split())]
heapq.heapify(heap)
for _ in range(M):
value = -heapq.heappop(heap)
value //= 2
heapq.heappush(heap, -value)
print(-sum(heap))コードの正しさ
- ヒープの最小キーは、負号を戻したときに現在価格の最大値です。したがって
heappopで、割引券を使うべき商品を取り出せます。 - 取り出した価格を整数除算で2分の1にし、負号を戻して同じヒープへ入れ直します。各反復の後も、ヒープは全商品の現在価格を一つずつ表します。
- 一枚ごとに最大価格を選ぶ交換が成り立つので、
M回後の価格合計が最小になります。負のキーの合計にもう一度負号を付ければ、実際の合計価格です。
ヒープ化はO(N)、割引券一枚ごとの取り出しと追加はO(log N)なので、時間計算量はO(N + M log N)、追加メモリはO(N)です。MがNより大きくても同じ商品へ何度も使えます。価格が1なら0になり、価格0へ残りの券を使っても0のままです。公式サンプル3のように一つの商品を多数回半分にする場合も、特別な分岐は要りません。
ABC212 D：共通offsetで全ボールの加算を遅延する
ABC212 Dの公式問題文の操作1は値Xを袋へ追加し、操作2は袋の全ボールへXを加え、操作3は最小のボールを取り出します。Q <= 2 * 10^5、1 <= X <= 10^9で、操作3の直前に袋が空でないことが保証されています。公式EditorialのD – Querying Multisetは、全体への加算を累積値として外へ出す方法を示しています。
実値をbaseとoffsetに分ける
ヒープに入れる各値をbase、全ボールに共通する加算をoffsetとし、実際の値をbase + offsetと表します。操作1で値xを追加するとき、今のoffsetを含めてxになるようにx - offsetをヒープへ入れます。操作2は全要素を書き換えず、offsetだけを増やします。同じ値を全要素へ加えても大小関係は変わらないため、操作3では最小のbaseを取り出してoffsetを戻せば、最小の実値になります。
公式サンプル1の操作列は、1 3、1 5、3、2 2、3です。状態を実値と分けて追うと、最初の二つの追加後はbase=[3, 5]、offset=0です。最初の操作3で3を取り出した後、操作2でoffset=2になり、残ったbase=5の実値は7になります。出力は3、7です。
import heapq
Q = int(input())
heap = []
offset = 0
answer = []
for _ in range(Q):
query = list(map(int, input().split()))
if query[0] == 1:
heapq.heappush(heap, query[1] - offset)
elif query[0] == 2:
offset += query[1]
else:
answer.append(heapq.heappop(heap) + offset)
for value in answer:
print(value)コードの正しさ
- 追加操作では、ヒープのキーへ現在の
offsetを引くため、保存したbaseにoffsetを足すと追加時の実値xへ戻ります。 - 全体加算では
offsetだけを更新します。すべてのボールへ同じ値を足すので、実値の大小順とbaseの大小順は一致したままです。 - 取り出し操作は最小の
baseを一つ取り出し、offsetを足して出力します。操作3の直前は袋が空でないため、取り出し前の空判定は必要ありません。
追加と取り出しはO(log Q)、全体加算はO(1)です。全体でO(Q log Q)時間、ヒープと出力列でO(Q)メモリを使います。offsetが大きくなって保存したbaseが負になっても、それは共通加算を取り除いたキーなので問題ありません。公式サンプル2の出力5000000000のように32ビット整数を超える値も、Pythonの整数でそのまま扱えます。入力のクエリ種別は問題文の1、2、3であり、コードのquery[0]はその種別、query[1]は2番目の入力値を指す0始まりの添字です。
ABC137 D：締切の厳しい枠から後ろ向きに候補を追加する
ABC137 Dの公式問題文では、仕事jをするとA_j日後に報酬B_jを受け取ります。jは1始まりの仕事番号、A_jは仕事jをしてから報酬を受け取るまでの日数です。一日に一件まで働き、今日からM日後までに受け取れる報酬の合計を最大化します。制約はN, M <= 10^5、A_j <= 10^5、B_j <= 10^4です。公式EditorialのD – Summer Vacationの見方に合わせ、最後の枠から前へ戻ります。
後ろから数えると候補条件が単純になる
最後に働ける枠を、後ろから数える1始まりの番号k=1と置き、そこから一つ前の枠へ移るたびにkを1増やします。仕事番号はjと分けます。仕事jのA_jは報酬を受け取るまでの日数なので、後ろからk番目の枠で間に合う候補条件はA_j <= kです。したがって、k=1ではA_j=1の仕事だけを候補にし、k=2ではA_j=2の仕事を追加して、候補中の最大報酬を選びます。掲載コードのループ変数iは、この後ろからの枠番号kに対応します。コードは仕事番号jを使って並べ替えず、A_jをキーに報酬B_jをby_daysへまとめます。
公式サンプル1はM=4で、仕事は(A_j,B_j)=(4,3),(4,1),(2,2)です。k=1では候補がなく、k=2でA_j=2の報酬2を追加して選びます。k=3では候補が増えず、k=4でA_j=4の報酬3と1を追加して最大の3を選びます。合計は2+3=5です。A_j > Mの仕事は、どの枠に置いても受け取りが間に合わないので最初から除きます。
import heapq
N, M = map(int, input().split())
by_days = [[] for _ in range(M + 1)]
for _ in range(N):
a, b = map(int, input().split())
if a <= M:
by_days[a].append(b)
heap = []
answer = 0
for i in range(1, M + 1):
for reward in by_days[i]:
heapq.heappush(heap, -reward)
if heap:
answer -= heapq.heappop(heap)
print(answer)コードの正しさ
kまで候補を広げた時点で、ヒープにはA_j <= kを満たし、まだ使っていない仕事jの報酬だけが入ります。A_j > Mの仕事を除く条件も、受け取り期限を満たすために必要です。- 現在の枠より前の枠ほど候補を広く選べるため、現在の候補の中で報酬最大の仕事を選ぶ交換ができます。現在の枠で小さい報酬を選び、より大きい候補を残す解があれば、二つを交換しても後の枠の締切条件を壊しません。
- 取り出した仕事はヒープから消えるので、一つの仕事を二度使いません。各枠で最大報酬を積み上げる結果が、M個の枠で得られる最大合計になります。
ループ自体はM回です。仕事の追加は最大N回、取り出しも各仕事一回までなので、時間計算量はO(M + N log N)、by_daysとヒープを合わせた追加メモリはO(M + N)です。候補が空の枠は働かず、そのまま前の枠へ進みます。公式サンプル3のM=1、A_j=2の仕事は候補が空なので出力0です。ここでもjは1始まりの仕事番号、A_jはその仕事の報酬受取までの日数です。後ろからの枠番号はk、掲載コードのiはkに対応し、両者を混同しません。Pythonのby_days[a]は仕事番号ではなく、日数aをキーにした配列です。
ABC325 D：開始時刻から前向きに掃引し、期限切れを捨てる
ABC325 Dの公式問題文では、商品iが時刻T_iに印字機の範囲へ入り、D_iマイクロ秒後の時刻T_i + D_iに出ます。商品は入る瞬間と出る瞬間にも印字でき、印字間隔は1マイクロ秒です。制約は1 <= N <= 2 * 10^5、1 <= T_i, D_i <= 10^18です。公式EditorialのD – Printing Machineに沿って、開始時刻順に追加し、退出時刻の最小値を選びます。
区間の端を含めて期限を管理する
商品iの印字可能な時間は、整数時刻の閉区間[T_i, T_i + D_i]です。nowまでに入った商品を退出時刻end=T_i+D_iの最小ヒープへ追加し、end < nowだけを期限切れとして捨てます。end == nowはまだ印字できるため、等号を含めて削除してはいけません。ヒープが空なら、次に商品が入る時刻へnowをジャンプします。候補があれば、最も早く退出する商品を印字して、次の印字時刻をnow+1にします。
公式サンプル1の5商品では、時刻1に退出時刻2、2、3、5の候補が入ります。時刻1で退出時刻2、時刻2でもう一つの2、時刻3で3、時刻4で5を選ぶと、印字数は4です。時刻4では退出時刻3の候補がすでに過ぎているため捨てます。開始時刻が同じ商品も同じ時刻に一括で追加し、退出時刻が同じ商品はヒープの同順位として扱えます。
import heapq
N = int(input())
products = []
for _ in range(N):
t, d = map(int, input().split())
products.append((t, t + d))
products.sort()
heap = []
index = 0
now = 0
answer = 0
while index < N or heap:
if not heap:
now = max(now, products[index][0])
while index < N and products[index][0] <= now:
_, end = products[index]
heapq.heappush(heap, end)
index += 1
while heap and heap[0] < now:
heapq.heappop(heap)
if heap:
heapq.heappop(heap)
answer += 1
now += 1
print(answer)コードの正しさ
- 商品を開始時刻でソートし、
products[index][0] <= nowのものをすべて追加するので、ヒープへの追加漏れがありません。追加済み商品のうちend < nowだけを除くため、残るヒープは現在時刻に印字できる未処理商品の集合です。 - 同じ
nowで別の商品を選ぶ解があれば、退出時刻が最も早い商品をその時刻に置き換えられます。早く退出する商品を先に処理しても、もう一方の商品を後で印字できる可能性は狭まりません。したがって最小のendを選ぶ貪欲法が成立します。 - ヒープが空のときは、次の商品の開始時刻より前に印字できる商品がありません。次の開始時刻へジャンプしても、印字できる機会を失いません。印字した後に時刻を1進めることで、印字間隔1マイクロ秒も保たれます。
商品を開始時刻でソートし、各商品をヒープへ一度追加し、一度だけ取り出すので、時間計算量はO(N log N)、追加メモリはO(N)です。Pythonのindexは0始まりですが、商品時刻T_iとD_iは問題文の値をそのまま使います。10^18級の時刻でも整数のまま扱え、空白時間を一時刻ずつ走査しません。公式サンプル2のように二つの時間帯が遠く離れていても、ジャンプで処理できます。
四問を選び分けるチェック
問題文をコードへ写す前に、次の表の順で問い直します。
| 問題の中心 | ヒープのキー | 候補を更新する順序 | 境界の確認 |
|---|---|---|---|
| ABC141 Dの現在価格の最大値 | 価格の負数 | 割引一枚ごとに最大値を取り出して更新 | 価格0の後も券を使える |
| ABC212 Dの最小値と全体加算 | 追加値 - offset |
操作2ではoffsetだけを更新 | 取り出し時にoffsetを戻す |
| ABC137 Dの仕事の報酬と締切 | 報酬の負数 | 後ろからk=1..MでA_j <= kを追加 |
A_j > Mは候補外 |
| ABC325 Dの印字可能商品の最早退出 | T_i + D_i |
開始時刻順に追加し、時刻を前へ進める | end < nowだけを期限切れにする |
- 最小値か最大値かを決め、最大値なら負のキーへ変換します。
- 全要素への同じ加算があるなら、追加時にoffsetを引き、取り出し時にoffsetを足す不変条件を書きます。
- 候補が有効になる日や時刻を決めます。締切の日数は後ろから、開始・退出の区間は開始時刻から処理します。
- 区間の端が含まれるかを確認します。ABC325 Dでは
end == nowが有効なので、期限切れ判定はend < nowです。
中央値や上位Kのように一つの極値だけでは足りない問題では、この4問の単一ヒープをそのまま当てはめません。必要な順序統計量を言葉で定め、キー、追加時点、削除境界の三つが揃ったときに、対応するデータ構造を選びます。
