Pythonで理解する蟻本「2-4 食物連鎖（POJ 1182）」(p.85)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 食物連鎖（POJ 1182）」(p.85)
のコードをPythonで書き直したものとなっています。
入力
入力例
100 7
1 2 2 2 1 2 1
101 1 2 3 1 3 5
1 2 3 3 3 1 5
解答
# 入力（Tは情報のタイプ）
N, K = map(int,input().split())
T = list(map(int,input().split()))
X = list(map(int,input().split()))
Y = list(map(int,input().split()))
# Union-Find木のソースコード
class UnionFindTree():
# n要素で初期化
def __init__(self, n):
self.n = n
self.par = list(range(n))
self.rank = [0] * n
# 木の根を求める
def find(self, x):
if self.par[x] == x:
return x
else:
self.par[x] = self.find(self.par[x])
return self.par[x]
# xとyの属する集合を併合
def unite(self, x, y):
x = self.find(x)
y = self.find(y)
if x == y:
return
if self.rank[x] < self.rank[y]:
self.par[x] = y
else:
self.par[y] = x
if self.rank[x] == self.rank[y]:
self.rank[x] += 1
# xとyが同じ集合に属するか否か
def same(self, x, y):
return self.find(x) == self.find(y)
# Union-Find木を初期化
# x, x + N, x + 2 * N を x-A, x-B, x-Cの要素とする
UFT = UnionFindTree(N * 3)
ans = 0
for i in range(K):
t = T[i]
x = X[i] - 1 # 0, ..., N - 1の範囲にする
y = Y[i] - 1
# 正しくない番号
if x < 0 or N <= x or y < 0 or N <= y:
ans += 1
continue
if t == 1:
# 「xとyは同じ種類」という情報
if UFT.same(x, y + N) or UFT.same(x, y + 2 * N):
ans += 1
else:
UFT.unite(x, y)
UFT.unite(x + N, y + N)
UFT.unite(x + N * 2, y + N * 2)
else:
# 「xとyを食べる」という情報
if UFT.same(x, y) or UFT.same(x, y + 2 * N):
ans += 1
else:
UFT.unite(x, y + N)
UFT.unite(x + N, y + 2 * N)
UFT.unite(x + 2 * N, y)
print(ans)
Pythonで理解する蟻本「2-4 Union-Find木」(p.81)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 Union-Find木」(p.81)
のコードをPythonで書き直したものとなっています。
Union-Find木の実装
コード
MAX_N = 10 ** 8 par = [0] * MAX_N # 親 rank = [0] * MAX_N # 木の深さ # n要素で初期化 def init(n): for i in range(n): par[i] = i rank[i] = 0 # 木の根を求める def find(x, par): if par[x] == x: return x else: par[x] = find(par[x], par) return par[x] # xとyの属する集合を併合 def unite(x, y, par, rank): x = find(x, par) y = find(y, par) if x == y: return if rank[x] < rank[y]: par[x] = y else: par[y] = x if rank[x] == rank[y]: rank[x] += 1 # xとyが同じ集合に属するか否か def same(x, y, par, rank): return find(x, par) == find(y, par)
実行例
コード
init(100) # 要素数が100のUnion-Find木を作る print(same(0, 1, par, rank)) # 0と1が同じ集合に属するかを確認する unite(0, 1, par, rank) # 0と1の属する集合を併合する print(same(0, 1, par, rank)) # 0と1が同じ集合に属するかを確認する
実行結果
False True
classを使ったUFTの実装
上記の実行例を見ると、関数による実装では少し扱いづらいことが分かります。
そのため、ここではclassによる実装も載せておこうと思います。
（実際にこのコードを使ってUnion-Find木の問題を解けることは確認していますが、もしバグを発見された場合は僕のTwitterアカウント（クルトン (@kuruton456) | Twitter）に連絡して頂けると助かります。）
コード
class UnionFindTree(): # n要素で初期化 def __init__(self, n): self.n = n self.par = list(range(n)) self.rank = [0] * n # 木の根を求める def find(self, x): if self.par[x] == x: return x else: self.par[x] = self.find(self.par[x]) return self.par[x] # xとyの属する集合を併合 def unite(self, x, y): x = self.find(x) y = self.find(y) if x == y: return if self.rank[x] < self.rank[y]: self.par[x] = y else: self.par[y] = x if self.rank[x] == self.rank[y]: self.rank[x] += 1 # xとyが同じ集合に属するか否か def same(self, x, y): return self.find(x) == self.find(y)
実行例
コード
UFT = UnionFindTree(100) # 要素数が100のUnion-Find木を作る print(UFT.same(0, 1)) # 0と1が同じ集合に属するかを確認する UFT.unite(0, 1) # 0と1の属する集合を併合する print(UFT.same(0, 1)) # 0と1が同じ集合に属するかを確認する
実行結果
False True
Pythonで理解する蟻本「2-4 二分探索木」(p.75)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 二分探索木」(p.75)
のコードをPythonで書き直したものとなっています。
二分探索木の実装
コード
# ノードを表す構造体 class node: def __init__(self): self.val = None self.lch = None self.rch = None # 数xを追加 def insert(p, x): if p == None: q = node() q.val = x return q else: if x < p.val: p.lch = insert(p.lch, x) else: p.rch = insert(p.rch, x) return p # 数xを検索 def find(p, x): if p == None: return False elif x == p.val: return True elif x < p.val: return find(p.lch, x) else: return find(p.rch, x) # 数xを削除 def remove(p, x): if p == None: return None elif x < p.val: p.lch = remove(p.lch, x) elif x > p.val: p.rch = remove(p.rch, x) elif p.lch == None: q = p.rch del p return q elif p.lch.rch == None: q = p.lch q.rch = p.rch del p return q else: q = node() q = p.lch while q.rch.rch != None: q = q.rch r = q.rch q.rch = r.lch r.lch = p.lch r.rch = p.rch del p return r return p
実行例
コード
root = None root = insert(root, 1) # 1を追加 print(find(root, 1)) # 1を検索 root = remove(root, 1) # 1を削除 print(find(root, 1)) # 1を検索
実行結果
True False
Pythonで理解する蟻本「Fence Repair（PKU 3253）」(p.75)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「Fence Repair（PKU 3253）」(p.75)
のコードをPythonで書き直したものとなっています。
入力
入力例
3
8 5 8
解答
# 順位キュー（優先度付きキュー）をインポートする
import heapq
# 入力
N = int(input())
L = list(map(int,input().split()))
ans = 0
# 順位キュー（優先度付きキュー）を用意
que = []
for i in range(N):
heapq.heappush(que, L[i])
# 坂が1本になるまで適用
while len(que) > 1:
# 一番短い板, 次に短い板を取り出す
l1 = heapq.heappop(que)
l2 = heapq.heappop(que)
# それらを併合
ans += l1 + l2
heapq.heappush(que, l1 + l2)
print(ans)
Pythonで理解する蟻本「2-4 Expedition（POJ 2431）」(p.73)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 Expedition（POJ 2431）」(p.73)
のコードをPythonで書き直したものとなっています。
入力
入力例
C++のpriority_queueは最大値を返しますが、Pythonのheapqは最小値を返します。 そのため、下のコードではヒープに追加する前に-1を掛け、取り出した値にもう一度-1を掛けることで最大値を得ています。
25 10 4
10 14 20 21
10 5 2 4
解答
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)
Pythonで理解する蟻本「2-4 プライオリティキューとヒープ」(p.71)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-4 プライオリティキューとヒープ」(p.71)
のコードをPythonで書き直したものとなっています。
ヒープの実装例
MAX_N = 10 ** 8 heap = [0] * MAX_N sz = 0 def push(x): global sz # 自分のノード番号 i = sz sz += 1 while i > 0: # 親のノードの番号 p = (i - 1) // 2 # もう逆転していないなら抜ける if heap[p] <= x: break # 親のノードの数字を下ろして、自分は上に heap[i] = heap[p] i = p heap[i] = x def pop(): global sz # 最小値 ret = heap[0] # 根に持ってくる値 sz -= 1 x = heap[sz] # 根からおろしていく i = 0 while i * 2 + 1 < sz: # 子同士を比較 a = i * 2 + 1 b = i * 2 + 2 if b < sz and heap[b] < heap[a]: a = b # もう逆転していないなら終わり if heap[a] >= x: break # 子の数字を持ち上げる heap[i] = heap[a] i = a heap[i] = x return ret
PythonとC++の違い
Pythonではheapqが、C++のpriority_queueに相当します。
しかし完全に同じというわけではなく、C++のpriority_queueが最大値を返す一方で、Pythonのheapqは最小値を返します。
またC++においては、popはコンテナの先頭側の要素を削除するだけで、返り値はありません。
一方でPythonにおいては、heappopはコンテナの先頭側の要素を削除し、その値を返します（返り値がある）。
そのため、蟻本ではスタックの先頭側の要素を出力した後にその値を削除していますが、この記事のコードでは先頭側の要素を削除する際に値を変数に受け取り、その値を出力しています。
（蟻本のコードは「出力→値の削除」の順番になっているが、この記事のコードでは「値の削除→出力」の順番になっている。）
標準ライブラリを使ったコード
コード
import heapq # 宣言 pque = [] # 要素の追加 heapq.heappush(pque, 3) heapq.heappush(pque, 5) heapq.heappush(pque, 1) # 空になるまでループ while len(pque) != 0: # 最小値の取得および削除 MIN = heapq.heappop(pque) # MINに最小値を代入し削除する print(MIN)
実行結果
1 3 5
最大値を取り出す目的でheapqを使う
heapqには最大値を取り出すためのメソッドがありません。
そのため、最大値を取り出す目的でheapqを使いたいときは、要素に-1を掛けてヒープに追加し、取り出した後にもう一度-1を掛けるとよいでしょう。
コード
import heapq # 宣言 pque = [] # 要素の追加 heapq.heappush(pque, -1 * 3) # 要素を追加する際に-1を掛ける heapq.heappush(pque, -1 * 5) heapq.heappush(pque, -1 * 1) # 空になるまでループ while len(pque) != 0: # 最小値の取得および削除 MAX = -1 * heapq.heappop(pque) # 受け取る際にもう一度-1を掛ける print(MIN)
実行結果
5 3 1
Pythonで理解する蟻本「2-3 重複組み合わせ」(p.67)
この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「2-3 重複組み合わせ」(p.67)
のコードをPythonで書き直したものとなっています。
入力
入力例
3 3
1 2 3
10000
解答
MAX_N = 1000
MAX_M = 1000
# 入力
n, m = map(int,input().split())
a = list(map(int,input().split()))
M = int(input())
dp = [[0] * (MAX_M + 1) for _ in range(MAX_N + 1)] # DPテーブル
# 1つも選ばない方法は常に一通り
for i in range(n + 1):
dp[i][0] = 1
for i in range(n):
for j in range(1, m + 1):
if j - 1 - a[i] >= 0:
# Pythonでは負の数でも直接余りを求められる（Mを足さなくてもよい）
dp[i + 1][j] = (dp[i + 1][j - 1] + dp[i][j] - dp[i][j - 1 - a[i]]) % M
else:
dp[i + 1][j] = (dp[i + 1][j - 1] + dp[i][j]) % M
print(dp[n][m])
