蟻本をpythonで: Sec. 2-4, データ構造

page 73: Expedition (POJ 2421)

# -*- coding: utf-8 -*-

from heapq import *

N = int(raw_input())
L = int(raw_input())
P = int(raw_input())
A = map(int,raw_input().split())
B = map(int,raw_input().split())

h = []

fuel = P
position = 0
ans = 0

A.append(L)
B.append(0)
N += 1

def solve():

    h = []

    fuel = P
    position = 0
    ans = 0


    for i in range(N):
        dl = A[i]-position

        while fuel - dl <0:
            if len(h) == 0:
                return -1
            fuel += -heappop(h)
            ans += 1
        fuel -=dl
        position = A[i]
        heappush(h,-B[i])

    return ans

print solve()

C++のSTLと異なり、pythonのheappop()は最小要素から出てくるので、符号を反転させて対応。

page75: Fence Repair (PKU 3253) 再び

# -*- coding: utf-8 -*-

from heapq import *

N = int(raw_input())
L = map(int,raw_input().split())

h = []

cost = 0

for i in L:
    heappush(h,i)

while len(h)>1:
    l1 = heappop(h)
    l2 = heappop(h)
    cost += l1 + l2
    heappush(h,l1+l2)

print cost

heapqによる解法。
最小要素から持ってくるのでpythonのheapqそのまま。計算量がO(L, n log n)。

page85: 食物連鎖 (POJ 1182)

# -*- coding: utf-8 -*-

from heapq import *
from unionfind import *

N = int(raw_input())
K = int(raw_input())
information = []
while 1:
    a = raw_input()
    if a == "":
        break
    information.append(map(int,a.split()))

u = unionfind(3 * N) 

ans = 0
for i in information:
    infclass = i[0]
    x = i[1]
    y = i[2]

    if x < 0 or x >=N or y <0 or y >=N or infclass <1 or infclass >2:
        ans +=1
        continue
    if infclass == 1:
        if u.issame(x, y+N) or u.issame(x, y+2*N):
            ans +=1
            continue
        else:
            u.unite(x, y)
            u.unite(x+N, y+N)
            u.unite(x+2*N, y+2*N)

    if infclass ==2:
        if u.issame(x, y) or u.issame(x, y+2*N):
            ans +=1
            continue
        else:
            u.unite(x, y+N)
            u.unite(x+N, y+2*N)
            u.unite(x+2*N, y)

print ans

いやはや面白い。
あえて要素を複製することでグループ分けの問題に落とし込めるので、Union-Findによって高速に処理できる。

Sec2-4を終えて

プライオリティキュー、二分探索木、Union-Find木の導入と練習問題。

とりあえずheapqとunionfindのパッケージを使って書いてみた。

Sec2-5に行く前にそれぞれ自分で実装してみようと思う。