Arantium Maestum

プログラミング、囲碁、読書の話題

蟻本初級編攻略 - 2-4 Expedition

動的計画法関連のところはPythonでやるとTLEする問題もいくつかあり、OCamlでやることも視野に入れつつ少し寝かせておく。

というわけで先にPriority Queueを使った問題3つ。まずは蟻本に載っているPOJ問題から:

POJ - Expedition

2431 -- Expedition

現在地から目的地への直線上にある複数のガソリンスタンドに最小で何回止まって給油する必要があるか、という問題。

ガソリンが足りなくなる限界までの区間にあるガソリンスタンドの給油量をPriorityQueueにいれて最大値を取り出し、ガソリンが足りなくなる限界距離を更新して考慮に入れるスタンドをQueueにいれて最大値を取り出し、というループを限界距離が目的距離に到達するまで続ける:

from queue import PriorityQueue

n = int(input())
xs = [tuple(int(x) for x in input().split()) for _ in range(n)]
l, p = map(int, input().split())

current = p
count = 0
q = PriorityQueue()
xs = list(sorted([(l-x, v) for x, v in xs], reverse=True))

while current < l:
    while xs and xs[-1][0] <= current:
        q.put(-xs.pop()[1])
    if not q:
        break
    current -= q.get()
    count += 1

print(count if current >= l else -1)

PythonのPriorityQueueやheapqはmin-heapしか提供していない。ちょっとショック。なので「最大値を常にとってくる」という挙動のためにマイナスでかけている。その点以外は比較的素直な実装ではないだろうか。

それではAtCoderの類題2問を解いてみる。

Code Thanks Festival 2017 C問題 - Factory

code-thanks-festival-2017-open.contest.atcoder.jp

code-thanks-festival-2017-open.contest.atcoder.jp

プレゼントを作るごとにスピードが劣化していく複数の機械をどう使うと最小の時間で一定数のプレゼントが作れるか、という問題。

キューに「かかる時間」と「劣化の度合い」をタプルとして入れて、かかる時間が最小のものを選んで、そして「かかる時間+劣化」と「劣化の度合い」のタプルをキューに戻して、とループしていく。ほしいプレゼント数の分だけループが回って「かかる時間」の和が答え。

import heapq
 
n, k = map(int, input().split())
xys = [tuple(int(x) for x in input().split()) for _ in range(n)]
 
heapq.heapify(xys)
 
time = 0
for _ in range(k):
    time += xys[0][0]
    heapq.heappushpop(xys, (xys[0][0]+xys[0][1], xys[0][1]))
 
print(time)

今回はheapqで実装してみた。heapqには普通のlistをヒープ化したり、ヒープ化されたリストに対するヒープ操作をしたり、という関数が用意されている。実装がリスト準拠なので、最小の要素をpeekするだけならxs[0]でできる。queueモジュールのPriorityQueueは実はスレッドセーフな実装になっており、その分のオーバーヘッドがかかるのでheapqのほうが定数倍効率がいい。

heapqがO(N)、heappushpopがO(logN)なので全部でO(n + k logn)の計算量。

ARC074 D問題 - 3N Numbers

arc074.contest.atcoder.jp

arc074.contest.atcoder.jp

3*N個の正の整数のリストからN個要素を取り除き、「前半N個の和 - 後半N個の和」を最大化する問題。

制約から見えてきたいくつかのポイント:

  • 前半の和を最大化し、後半の和を最小化する必要がある
  • 前半はA[:2*n]の要素のうちのN個によって構成される
  • 後半はA[n:]の要素のうちのN個によって構成される
  • 前半が使える部分と後半が使える部分のオーバーラップであるA[n:2*n]のどこに線を引くかがポイント

というわけで

  • i = 0 ~ nの各iごとにA[:n+i]からN個選んだ場合の和の最大値を求める
  • j = 0 ~ nの各ごとにA[2*n-j:]からN個選んだ場合の和の最小値を求める

ということをPriorityQueueを使って一歩ずつ算出していく。その上でi+j = nで差が最大になるiとjのペアを見つける。

from heapq import heapify, heappushpop
 
n = int(input())
xs = [int(x) for x in input().split()]
 
xs1 = xs[:n]
heapify(xs1)
r1 = [sum(xs1)]
for x in xs[n:2*n]:
    r1.append(r1[-1] + x - heappushpop(xs1, x))
 
xs2 = [-x for x in xs[2*n:]]
heapify(xs2)
r2 = [sum(xs2)]
for x in reversed(xs[n:2*n]):
    r2.append(r2[-1] + (-x) - heappushpop(xs2, -x))
 
print(max(a+b for a,b in zip(r1, reversed(r2))))

前半の和の最大値を求めるのに、A[n+i]の要素をキューにいれてから最小値を出して、(A[n+i] - 最小値)をそれまでの和に足している。後半も全く同じロジック、ただし「最大値をキューから取り出す」という挙動のためにキューに入れる・キューから取り出す値をすべてマイナスでかけている。

ほとんどの部分がO(N)だが、ループしている部分だけN回ループが回ってO(logN)のheappushpopが毎回行われているので、そこが計算量O(N logN)で大きい。

Priority Queue感想

やっぱり便利。今回の問題はみんなPriority Queueが前面に出てきているが、そうでなくても解法の一部として計算量を下げるのに活躍してくれる場面は多そう。

しかしmax heapがないのはやっぱり不便だな、微妙なところで実装を間違えてバグりそうだから要注意だ。