Pythonで理解する蟻本「2-4 Expedition（POJ 2431）」(p.73) - クルトンのプログラミング教室

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 Expedition（POJ 2431）」(p.73)
のコードをPythonで書き直したものとなっています。

入力
入力例
解答

入力

$L\,P\,N\\A_1\,…\,A_N\\B_1\,…\,B_N$

入力例

25 10 4
10 14 20 21
10 5 2 4

解答

C++のpriority_queueは最大値を返しますが、Pythonのheapqは最小値を返します。

そのため、下のコードではヒープに追加する前に-1を掛け、取り出した値にもう一度-1を掛けることで最大値を得ています。

import heapq
import sys

MAX_N = 10000

# 入力
L, P, N = map(int,input().split())
A = list(map(int,input().split()))
B = list(map(int,input().split()))

# 簡単なためゴールをガソリンスタンドに追加
A += [L]
B += [0]
N += 1

# ガソリンスタンドを管理する順位キュー
que = []

# ans：補給回数, pos：今いる場所, tank：タンクのガソリンの量
ans = 0
pos = 0
tank = P

for i in range(N):
    # 次に進む距離
    d = A[i] - pos
    
    # 十分な量になるまでガソリンを補給
    while tank - d < 0:
        if len(que) == 0:
            print("-1")
            sys.exit()
        tank += -1 * heapq.heappop(que)
        ans += 1
    
    tank -= d
    pos = A[i]
    heapq.heappush(que, -1 * B[i])

print(ans)