この記事では、螺旋本・AOJのALDS1コース(アルゴリズムとデータ構造入門)のトピック#7〜#12にPythonで解答を与えます。
前回の記事Pythonによる螺旋本・AOJの解答(ALDS1 #1~#6)の続きです。螺旋本・AOJとは何かについてはこちらの記事を参照して下さい。
目次
解答
トピック #7 木構造
ALDS1_7_A: Rooted Trees
class Node:
def __init__(self, pa=-1, chs=None):
self.pa = pa
self.chs = chs
n = int(input())
tree = {id:Node() for id in range(n)}
for _ in range(n):
id, _, *chs = map(int, input().split())
tree[id].chs = chs
for ch in chs:
tree[ch].pa = id
def set_depths(id, depth):
tree[id].depth = depth
for ch in tree[id].chs:
set_depths(ch, depth+1)
for id in tree:
if tree[id].pa == -1:
set_depths(id, 0)
break
for id, node in tree.items():
kind = "root" if node.pa == -1 else "internal node" if node.chs else "leaf"
print(f"node {id}: parent = {node.pa}, depth = {node.depth}, {kind}, {node.chs}")
-
親
pa
のid
への参照と子達chs
のid
への参照を保持するクラスNode
として節点を表現している。子全員への参照の代わりに、最左の子への参照と最右の兄弟ヘの参照だけで節点を表現し、任意の木を実質的に二分木として扱える"左-子 右-兄弟表現 (left-child, right-sibling representaiton)"という表現もあるけれど、子全員への参照を使った表現のほうがコードは読みやすいと思う -
id
をkeyとし対応するNode
のインスタンスをvalueとする辞書として木を表現している -
親が無であるときに親を-1と表示する問題なので、形式的に-1を無のidと扱っている
ALDS1_7_B: Binary Trees
class Node:
def __init__(self, chs, pa=-1):
self.chs = chs
self.pa = pa
n = int(input())
tree = {id:Node(chs=[-1,-1]) for id in range(-1, n)}
for _ in range(n):
id, *chs = map(int, input().split())
tree[id].chs = chs
for ch in chs:
tree[ch].pa = id
def set_d_and_h(id, d):
tree[id].d = d
h = 0
for ch in tree[id].chs:
if ch != -1:
h = max(h, set_d_and_h(ch, d+1) + 1)
tree[id].h = h
return h
for id in tree:
if tree[id].pa == -1:
set_d_and_h(id, 0)
break
for id, node in tree.items():
if id != -1:
sib = tree[node.pa].chs[1] if tree[node.pa].chs[0] == id else tree[node.pa].chs[0]
deg = 2 - node.chs.count(-1)
kind = "root" if node.pa == -1 else "internal node" if deg else "leaf"
print(f"node {id}: parent = {node.pa}, sibling = {sib}, degree = {deg}, depth = {node.d}, height = {node.h}, {kind}")
-
子が無のとき-1と表現されて入力されるので、その場合でも子についてのループを取れるよう、形式的に
id
が-1に対応する節点もいわゆる番兵として木に追加している。id
が-1の節点の親子情報が使われることはない -
深さと高さはそれぞれ帰納的に計算でき、うまくやればこのように同時に計算できる
ALDS1_7_C: Tree Walk
class Node:
def __init__(self, chs, pa=-1):
self.chs = chs
self.pa = pa
n = int(input())
tree = {id:Node(chs=[-1,-1]) for id in range(-1, n)}
for _ in range(n):
id, *chs = map(int, input().split())
tree[id].chs = chs
for ch in chs:
tree[ch].pa = id
def walk(id, order):
walked = ""
if id != -1:
if order == "Pre":
walked += f" {id}"
walked += walk(tree[id].chs[0], order)
if order == "In":
walked += f" {id}"
walked += walk(tree[id].chs[1], order)
if order == "Post":
walked += f" {id}"
return walked
for id in tree:
if tree[id].pa == -1:
for order in ["Pre", "In", "Post"]:
print(order + "order\n" + walk(id, order))
break
-
def walk
より上の部分は前問7_Bと同じ
ALDS1_7_D: Reconstruction of a Tree
_ = int(input())
pr_walk = [*map(int, input().split())]
in_walk = [*map(int, input().split())]
walked = []
def recon_tree(pr_walk, in_walk):
if pr_walk:
root = pr_walk[0]
i = in_walk.index(root)
recon_tree(pr_walk[1:i+1], in_walk[:i])
recon_tree(pr_walk[i+1:], in_walk[i+1:])
walked.append(root)
recon_tree(pr_walk, in_walk)
print(*walked)
- Post order walkを表示するだけであれば木を再構成するプロセス中にできて、再構成した木を使う必要が無いので、節点のインスタンスを作る必要が無い
トピック #8 二分探索木
連結リストでは挿入・(位置を指定した)削除は$O(1)$でできるものの、検索に$O(n)$かかってしまい、キーを指定した削除も検索が必要なので$O(n)$かかっていた。一方、二分探索木は、木の高さを$h$として挿入・検索・削除を全て$O(h)$で行えるデータ構造。$h$は最悪ケースでは$O(n)$だけど、ランダムに要素を挿入して木を構成するとき平均的には$O(\log n)$でとても高速。また、中間順巡回 (in-order traversal) によって、$\Theta(n)$の時間で節点をキーのソート順に訪れることもできる。
ALDS1_8_A: Binary Search Tree I
二分探索木で要素を挿入する問題。以下のIIIの問題に完全に含まれているので略
ALDS1_8_B: Binary Search Tree II
二分探索木で要素を挿入・検索する問題。以下のIIIの問題に完全に含まれているので略
ALDS1_8_C: Binary Search Tree III
二分探索木で要素を挿入・検索・削除する問題
class Node:
def __init__(self, key, left=None, right=None):
self.key = key
self.left = left
self.right = right
def insert(key):
global root
if root:
ch = root
while ch:
pa, ch = ch, ch.left if key < ch.key else ch.right
if key < pa.key:
pa.left = Node(key)
else:
pa.right = Node(key)
else:
root = Node(key)
def find(key):
node = root
while node and node.key != key:
node = node.left if key < node.key else node.right
print("yes" if node else "no")
def delete(key):
global root
pa, node = None, root
while node.key != key:
pa, node = node, node.left if key < node.key else node.right
if node.left and node.right:
pa, to_del = node, node.right
while to_del.left:
pa, to_del = to_del, to_del.left
node.key = to_del.key
else:
to_del = node
ch = to_del.left or to_del.right
if not pa:
root = ch
elif pa.left == to_del:
pa.left = ch
else:
pa.right = ch
def walk(node, order):
walked = ""
if node:
if order == "Pre":
walked += f" {node.key}"
walked += walk(node.left, order)
if order == "In":
walked += f" {node.key}"
walked += walk(node.right, order)
return walked
def show():
for order in ["In", "Pre"]:
print(walk(root, order))
root = None
cmds = {"print":show, "insert":insert, "find":find, "delete":delete}
for _ in range(int(input())):
cmd_name, *key = input().split()
cmds[cmd_name](*map(int, key))
-
親や子が無いときは、親や子を形式的に
None
として扱っている -
挿入と検索は見たとおりで、根から葉へ下っていくだけで分かりやすいけど、削除は少しややこしい
-
削除したい節点の子が高々1つの場合は、単にその子(子が無い場合は
None
)を削除したい節点の親の子にすれば、いわば"子に跡を継がせれば"済む。削除したい節点に子が2つある場合がやっかいで、例えば右の子に継いでもらおうとしたとき、右の子に既に左の子がいた場合、左の子が2つになってしまう。そこで、他の節点に継いでもらうことを考える。継いだ後も各節点が中間順巡回でソート順に現れなければならないけれど、中間順巡回で直前か直後の節点を引っこ抜いてきて置き換えればこれは満たされる。ここでは直後の節点(右の部分木を左に辿り続けたときの終点)を引っこ抜いてきている。直前(右の子がない)や直後(左の子がない)の節点には子が高々1つしかないので、子が高々1つの場合の先ほどの方法で引っこ抜ける -
各節点に親節点の情報を持たせなくても何とかなる。挿入や検索では親節点の情報は必要なく、必要となるのは削除のときだけだけど、削除したい節点にたどり着くまで道を辿る途中で親の情報を保持しておける
-
実は、pythonでは
A or B
はA
が真ならA
を返し、A
が偽ならB
を返している
なお、コルメン『アルゴリズムイントロダクション 第3版』の12章の章末で指摘されているように、この削除方法では問題が起きる場合がある。この削除方法では、削除したい節点に子が2つある場合、本来削除したい節点のインスタンスが削除されるのではなく、本来削除したい節点のインスタンスに別の節点のキー情報がコピーされて、その別の節点のインスタンスが削除されることになる。そのため、コードの別の場所でキーではなくインスタンスへの参照を行っている場合に、意図した節点とは異なる節点を参照してしまいかねない。この問題を解決するには、より複雑になってしまうものの、delete
関数部分を次のようにキー情報のコピーを使わない形に書き換えればよい:
def trans(pa, node, new_node):
if not pa:
global root
root = new_node
elif pa.left == node:
pa.left = new_node
else:
pa.right = new_node
def delete(key):
pa, node = None, root
while node.key != key:
pa, node = node, node.left if key < node.key else node.right
if node.left and node.right:
pa_heir, heir = node, node.right
while heir.left:
pa_heir, heir = heir, heir.left
if pa_heir != node:
trans(pa_heir, heir, heir.right)
heir.right = node.right
trans(pa, node, heir)
heir.left = node.left
else:
trans(pa, node, node.left or node.right)
ALDS1_8_D: Treap
- 螺旋本で略されているので、いったん略
二分探索木は木の高さを$h$として挿入・検索・削除を全て$O(h)$で行えて、ランダムに要素を挿入して木を構成するとき平均的には$h$は$O(\log n)$でとても高速だけど、最悪の挿入順に対しては$O(n)$になってしまう。任意の挿入順に対して木の高さの期待値あるいはより強く最悪値を$O(\log n)$に抑えるよう工夫した二分探索木を平衡二分探索木という。とくに、挿入順を選べないオンラインアルゴリズムで役立つ。
-
ツリープは、二分探索木に新しい変数である優先度をランダムに割り当てて二分ヒープにもしたもの。実質的に挿入順をランダム化する効果があり、任意の挿入順に対して木の高さの期待値を$O(\log n)$に抑えられる。ただし最悪値は$O(n)$
-
より強く木の高さの最悪値を$O(\log n)$に抑える平衡二分探索木には赤黒木(二色木)などがあるけど、よりややこしい
トピック #9 ヒープ
ALDS1_9_A: Complete Binary Tree
N = int(input())
heap = input().split()
for i in range(N):
print(f"node {i+1}: key = {heap[i]}, ", end="")
print(f"parent key = {heap[(i-1)//2]}, " if i else "", end="")
print(f"left key = {heap[2*i+1]}, " if 2*i+1 < N else "", end="")
print(f"right key = {heap[2*i+2]}, " if 2*i+2 < N else "")
- このコードのようにヒープの添え字を0始まりにすると、添え字
i
の親、左の子、右の子の添え字はそれぞれ(i-1)//2
,2*i+1
,2*i+2
になる。一方、ヒープの0番目にNone
などのダミー要素を入れて形式的に添え字を1始まりにすることもでき、その場合は親、左の子、右の子の添え字はそれぞれi//2
,2*i
,2*i+1
になる
ALDS1_9_B: Maximum Heap
max_heapify
を自分で書くと、
def max_heapify(A, i):
triad = {j:A[j] if j < N else -float("inf") for j in [i, 2*i+1, 2*i+2]}
i_max = max(triad, key=triad.get)
if i != i_max:
A[i], A[i_max] = A[i_max], A[i]
max_heapify(A, i_max)
N = int(input())
heap = [*map(int, input().split())]
for i in reversed(range(N//2)):
max_heapify(heap, i)
print("", *heap)
-
max_heapify
では、指定された親節点A[i]
とその子A[2*i+1]
,A[2*i+2]
の三つ組を考える。存在しない子は$-\infty$と扱う。もし親節点が子節点より小さければ、親節点を最大の子節点と入れ替えてからその子節点に対して再帰的にmax_heapify
を行う -
max_heapify
は子を辿っていくので、計算量は完全二分木の高さのオーダー$O(\log N)$。max_heapify
を全ての節点に対して行うと、高さ$h$ $(1 \le h \le \log N)$の部分木は高々$N/2^h$個しかないので、全体の計算量は$O(N \sum_{h=1}^{\log N} h/2^h)=O(N)$。ここで、$\sum_{h=0}^\infty h/2^h = O(1)$を用いた
heapq
モジュールを使うと:
from heapq import heapify
_, heap = input(), [-int(x) for x in input().split()]
heapify(heap)
print("", *[-x for x in heap])
-
heapq
モジュールは最小ヒープを扱うので、最大ヒープを最小ヒープとして扱うために負号をつけている
ALDS1_9_C: Priority Queue
以下のようにextract
やinsert
を自分で書くと、TLEになってしまう:
def max_heapify(A, i):
triad = {j:A[j] if j < len(A) else -float("inf") for j in [i, 2*i+1, 2*i+2]}
i_max = max(triad, key=triad.get)
if i != i_max:
A[i], A[i_max] = A[i_max], A[i]
max_heapify(A, i_max)
def extract(A):
print(A[0])
A[0] = A[-1]
A.pop()
max_heapify(A, 0)
def increase_key(A, i, key):
A[i] = key
while i and A[(i-1)//2] < A[i]:
A[(i-1)//2], A[i] = A[i], A[(i-1)//2]
i = (i-1)//2
def insert(A, key):
A.append(None)
increase_key(A, len(A)-1, key)
heap = []
while True:
op, *key = input().split()
if op == "end":
break
if op == "extract":
extract(heap)
else:
insert(heap, int(*key))
どうやらよほど工夫しなければ通らないようなので、諦めてheapq
モジュールを使うと通る:
from heapq import heappop, heappush
heap = []
while True:
op, *key = input().split()
if op == "end":
break
if op == "extract":
print(-heappop(heap))
else:
heappush(heap, -int(*key))
トピック #10 動的計画法
動的計画法には、メモ化再帰を用いたトップダウン方式と、部分問題を明示的に小さい順に解いていくボトムアップ方式がある。
-
トップダウン方式では関数の呼び出しのオーバーヘッドが余計にかかるデメリットがあるものの、全ての部分問題を解く必要が無い場合に必要な部分問題だけを解くメリットがある
-
一方、ボトムアップ方式ではどの順にボトムアップするかを考える必要があるデメリットがあるものの、下のフィボナッチ数の問題のように、部分問題の計算結果を途中で破棄できる場合に空間計算量を節約できるメリットがある
ALDS1_10_A: Fibonacci Number
トップダウン方式では:
from functools import lru_cache
@lru_cache(maxsize=None)
def fib(n):
if n < 2:
return 1
return fib(n-1) + fib(n-2)
print(fib(int(input())))
ボトムアップ方式では:
x = y = 1
for _ in range(int(input())):
x, y = x + y, x
print(y)
ALDS1_10_B: Matrix Chain Multiplication
トップダウン方式では:
from functools import lru_cache
n = int(input())
p = [int(input().split()[0]) for _ in range(n-1)] + [*map(int, input().split())]
@lru_cache(maxsize=None)
def n_mul(i, j):
if j - i == 1:
return 0
return min(n_mul(i,k) + n_mul(k,j) + p[i]*p[k]*p[j] for k in range(i+1,j))
print(n_mul(0, n))
- 行列$A_k \in \text{M}(p_k,p_{k+1})$の積$\prod_{k=0}^{n-1} A_k$の最小計算回数を考える。
n_mul(i, j)
は部分問題$\prod_{k=i}^{j-1} A_k$を解く
ボトムアップ方式では:
n = int(input())
p = [int(input().split()[0]) for _ in range(n-1)] + [*map(int, input().split())]
n_mul = [[0]*(n+1) for _ in range(n+1)]
for l in range(2, n+1):
for i in range(n-l+1):
j = i + l
n_mul[i][j] = min(n_mul[i][k] + n_mul[k][j] + p[i]*p[k]*p[j] for k in range(i+1,j))
print(n_mul[0][n])
- 部分問題のサイズ$l=j-i$が小さい順に部分問題の計算結果
n_mul[i][j]
を埋めていく
ALDS1_10_C: Longest Common Subsequence
Pythonでは、以下の自然なやり方では、トップダウン方式でもボトムアップ方式でも、途中でTLEしてしまう。時間計算量はともに$O(q|X||Y|)$である。
トップダウン方式では:
import sys
from functools import lru_cache
sys.setrecursionlimit(100000)
for _ in range(int(input())):
x, y = input(), input()
@lru_cache(maxsize=None)
def lcs(i, j):
if i == 0 or j == 0:
return 0
if x[i-1] == y[j-1]:
return lcs(i-1, j-1) + 1
return max(lcs(i-1,j), lcs(i,j-1))
print(lcs(len(x), len(y)))
- 再帰が最大再帰処理回数を超えてしまうので、最大再帰処理回数を適当な大きめの値(100000)に再設定している
ボトムアップ形式では:
for _ in range(int(input())):
x, y = input(), input()
l = [[0] * (len(y)+1) for _ in range(len(x)+1)]
for i in range(len(x)):
for j in range(len(y)):
l[i+1][j+1] = l[i][j] + 1 if x[i] == y[j] else max(l[i][j+1], l[i+1][j])
print(l[-1][-1])
ボトムアップ形式ではi
についてのループで各i
に対してl[i+1][:]
を計算しているけれど、これは直前のi
での計算結果l[i][:]
しか使わない。そのため、l[i][:]
を一時的な変数l[:]
に格納してi
毎に更新し、過去のi
でのl[i][:]
の値を順次破棄していくことで、空間計算量を$O(|X||Y|)$から$O(|Y|)$に減らせる。これに伴い、時間計算量もオーダーは変わらないものの定数倍減り、TLEせずに通るようになる:
def lcs(X, Y):
l = [0] * (len(Y) + 1)
for x in X:
l_pre = l[:]
for j, y in enumerate(Y):
if x == y:
l[j+1] = l_pre[j] + 1
elif l[j+1] < l[j]:
l[j+1] = l[j]
return l[-1]
for _ in range(int(input())):
print(lcs(input(), input()))
このコードは、関数lcs
を定義しない下の形に書き直すと、同じ内容のはずなのにTLEしてしまう。なぜ関数を定義する形に書き直せば計算量が減るのか分からない。もし分かるかたがいたら、教えて頂けると嬉しいです。
for _ in range(int(input())):
X, Y = input(), input()
l = [0] * (len(Y) + 1)
for x in X:
l_pre = l[:]
for j, y in enumerate(Y):
if x == y:
l[j+1] = l_pre[j] + 1
elif l[j+1] < l[j]:
l[j+1] = l[j]
print(l[-1])
ALDS1_10_D: Optimal Binary Search Tree
- 螺旋本で略されているので、いったん略
トピック #11 グラフI
グラフの全頂点を辺に沿って探索していく主なグラフ探索アルゴリズムとして、訪問した頂点をスタックのようにLast In, First Outで探索していく深さ優先探索(Depth First Search; DFS)と、訪問した頂点をキューのようにFirst In, First Outで探索していく幅優先探索(Breadth First Search; BFS)がある。グラフの有向・無向に関わらず使える。
ALDS1_11_A: Graph
n = int(input())
A = [[0] * n for _ in range(n)]
for _ in range(n):
u, _, *V = map(int, input().split())
for v in V:
A[u-1][v-1] = 1
for a in A:
print(*a)
-
隣接リスト表現では$O(E+V)$の空間計算量で済むけれど、指定された2頂点間に辺があるかどうかを判定するのにリスト検索ぶんの時間計算量がかかる。一方、隣接行列表現では$O(V^2)$の空間計算量がかかってしまうものの、指定された2頂点間に辺があるかどうかを判定するのに$O(1)$の時間計算量しかかからない。
-
このコードでは隣接行列
A
を具体的に構成しているけれど、この問題では実は頂点は1から順にしか入力されないようで、わざわざ行列を作らずとも行ごとに入出力処理しても通る
ALDS1_11_B: Depth First Search
n = int(input())
G = [None] * n
for _ in range(n):
u, _, *V = map(int, input().split())
G[u-1] = [v-1 for v in V]
d, f = [None] * n, [None] * n
time = 0
def visit(u):
global time
time += 1
d[u] = time
for v in G[u]:
if not d[v]:
visit(v)
time += 1
f[u] = time
for u in range(n):
if not d[u]:
visit(u)
print(u+1, d[u], f[u])
-
スタックを使っても書けるけれど、このように再帰を使って書いたほうがコードが簡潔になる
-
未発見(下の3分類で言う白)の頂点を順次訪問していく。各頂点の訪問は、まずその頂点を発見(discover)し、その頂点に隣接する未発見(下の3分類で言う白)の頂点を全て訪問してから、完了(finish)する。発見時刻、完了時刻はそれぞれ
d
,f
に記録される -
各頂点の状態は以下の3種類
- 白(未発見):
d
もf
もNone
- 灰(発見済だけど未完了):
d
に発見時刻が入っているけれど、f
はNone
- 黒(完了済):
d
に発見時刻、f
に完了時刻が入っている
- 白(未発見):
-
このアルゴリズムの時間計算量は、グラフのスケールに関して線形$\Theta(V+E)$。なぜなら、節点
u
についてのループの時間計算量は、$\Theta(V)$とvisit
の計算量との和である。一方、visit
は各節点に対して1回ずつ呼び出され、その計算量はその節点から出る辺の数のオーダーであるため、全ての節点についての和は$\Theta(E)$
ALDS1_11_C: Breadth First Search
n = int(input())
G = [None] * n
for _ in range(n):
u, _, *V = map(int, input().split())
G[u-1] = [v-1 for v in V]
d = [0] + [-1] * (n-1)
q = [0]
while q:
v = q.pop(0)
for u in G[v]:
if d[u] == -1:
d[u] = d[v] + 1
q.append(u)
for v in range(n):
print(v+1, d[v])
- このアルゴリズムの時間計算量もグラフのスケールに関して線形$\Theta(V+E)$。なぜなら、キュー
q
についてのwhile
ループの時間計算量は、各節点が1回ずつキューに挿入されv
としてポップされるために生じる$\Theta(V)$と、内側のfor
文の計算量の和である。一方、各for
文の計算量はその節点v
から出る辺の数のオーダーであるため、全ての節点についての和は$\Theta(E)$
ALDS1_11_D: Connected Components
n, m = map(int, input().split())
G = [[] for _ in range(n)]
for _ in range(m):
s, t = map(int, input().split())
G[s].append(t)
G[t].append(s)
reps = [None] * n
for s in range(n):
if reps[s] is None:
reps[s] = s
q = [s]
while q:
v = q.pop(0)
for u in G[v]:
if reps[u] is None:
reps[u] = s
q.append(u)
for _ in range(int(input())):
s, t = map(int, input().split())
print("yes" if reps[s] == reps[t] else "no")
-
無向グラフの連結成分分解は、深さ優先探索か幅優先探索を用いて$O(V+E)$でできる。このコードでは幅優先探索を採用している。有向グラフの場合、一方通行では無くお互いに行き来できるかどうかで定義される強連結成分への分解はもう少しやっかい
-
各節点が属す連結成分の代表元を
reps
に記録しておくことによって、二つの節点が同じ連結成分に属すかどうかはreps
が同じかどうかを見れば$O(1)$で分かる。未訪問の節点はreps
がNone
。代表元としては、その連結成分で最初に訪れた節点s
を取っている
トピック #12 グラフII
ALDS1_12_A: Minimum Spanning Tree
from heapq import heapify, heappop
n = int(input())
G = [[*map(int, input().split())] for _ in range(n)]
heap = [[float("inf"), v] for v in range(n)]
heap[0] = [0, 0]
weights = 0
while heap:
heapify(heap)
w, u = heappop(heap)
weights += w
for i, [w, v] in enumerate(heap):
if 0 <= G[u][v] < w:
heap[i] = [G[u][v], v]
print(weights)
-
最小全域木は連結な無向重み付きグラフに対して定義され、貪欲法で求められることが知られている。辺の集合Aに対して、Aのどの辺とも交差しない任意のカットを取り、そのカットと交差する重み最小の辺をAに追加していけばよい。具体的には、1頂点だけの木からなる森から始めて木同士をくっつけて大きな木にしていくKruskalのアルゴリズムと、1頂点だけの木から始めて辺を加えて木を成長させていくPrimのアルゴリズムがある。このコードはPrimのアルゴリズムに基づいており、各頂点の追加コストを記録しながら、最も追加コストが少ない頂点から順に追加していく。追加する度に、追加された頂点から出る各辺が入る各頂点の追加コストを、もし小さく更新できれば更新する
-
これらのアルゴリズムの時間計算量はそれぞれのアルゴリズム中で用いられるデータ構造をどう実装するかに依存するが、グラフのサイズに関して対数線形で実現できる。Kruskalのアルゴリズムでは互いに素な集合族をどう実装するかに依り、ある実装では$O(E \log V)$を実現できる。Primのアルゴリズムでは優先度付きキューをどう実装するかに依り、このコードのように2分ヒープで実装すると同じく$O(E \log V)$を実現できる。これはフィボナッチヒープを用いると$O(E + V\log V)$に改善できる。連結なグラフでは$V-1 \le E$より$V=O(E)$なので、これは改善になっている。一般のグラフに対しては$V=\Omega(E^{1/2})$であり、$V$が$o(E)$であれば真の改善である
-
ただし、このコードでは隣接リスト表現を用いずに隣接行列表現を用いることで、ヒープから取り出された各要素に対してヒープにわたる
for
文を実行することになっており、ヒープのサイズが最大$V$であることを考えると、計算量に$O(V^2)$の寄与をしており計算量のオーダーは悪化していると思われる -
さらに、pythonの
heapq
モジュールではヒープ条件を保ったままある節点の値を減少させるいわゆるDecreaseKey機能が実装されていないため、仕方なく節点の値を減少させてからheapify
を毎回呼び出してヒープ条件を回復させているけれど、heapify
は$O(n)$なので$O(\log n)$のDecreaseKeyより遅い。今、ヒープのサイズ$n$は最大$V$であり、heapify
は$V$回呼び出されるため、heapify
から来る計算量の寄与も合計で$O(V^2)$である。一応、heapq
には非公式に_siftdown(heap, start_pos, pos)
という関数があり、これを用いればpos
にある節点の値を減少させたせいでヒープ条件を満たさなくなったヒープのヒープ条件を$O(\log n)$で回復させることができるが、わざわざ使わなくてもheapify
でも通るので使っていない -
各節点の追加コストからなる最小ヒープ
heap
を構成し、追加コストが最小の節点を取り出し、そのコストを重みの合計値weights
に加算していく。取り出した節点u
に隣接した節点たちv
の追加コストを更新するために、取り出した節点の追加コストw
だけではなくどの節点u
を取り出したかの情報も必要なので、ヒープには追加コストw
と節点番号u
のからなるリスト[w, u]
を入れておく。Pythonではリストの比較は最初の成分から順に辞書式で行うので、これでheappop
で正しく最小コストの節点[w, u]
を取り出せる -
取り出した節点
u
に隣接した(つまりG[u][v]
が非負な)各節点v
に対して、もしG[u][v]
がv
の追加コストw
より低ければ、v
の追加コストをw
からG[u][v]
に更新する -
AOJでは使えないけれど、AtCoderなどscipyが使える環境下では、自分でコードを書かなくても関数
scipy.sparse.csgraph.minimum_spanning_tree()
を使えば最小全域木が求まる
ALDS1_12_B: Single Source Shortest Path I
ALDS1_12_C: Single Source Shortest Path II
この二つの問題は、入力のサイズと制限時間が違うだけで内容は同じ。単一始点最短経路問題は重み付きの有向グラフに対して定義され、始点から到達可能な負閉路(重みの和が負となる閉路)が存在する場合に解が無くなる。
-
一般の重み付き有向グラフに対しては、$O(VE)$のBellman-Fordアルゴリズムで解くことができる。このアルゴリズムは始点から到達可能な負閉路の存在を検出できる
-
グラフが非巡回つまりDAGであると分かっていれば、負閉路は存在しない。この場合、トポロジカルソートを用いることで計算量を$O(V+E)$に落とすことができる
-
一方、重みが全て非負と分かっていても、負閉路は存在しない。この場合、Dijkstraのアルゴリズムで解くことができる。このアルゴリズムの計算量は優先度付きキューの実装によっていて、2分ヒープを用いると$O((V+E)\log V)$、フィボナッチヒープを用いると$O(V \log V + E)$である
この問題では非巡回かは分からないけれど重みが全て非負とは分かっているので、Dijkstraのアルゴリズムで解く。Dijkstraのアルゴリズムは先の最小全域木のPrimのアルゴリズムと類似していて、始点から最も距離$d$が近い頂点から順に最短経路およびその距離$d$を確定させていく。確定する度に、確定された頂点から出る各辺が入る各頂点の距離$d$を、もし短く更新できれば更新する。
先ほどの最小全域木の問題のPrimのアルゴリズムのコードとよく似た次のコードでは、Iは解けるけれどIIはTLEになる:
from heapq import heapify, heappop
n = int(input())
G = [None] * n
for _ in range(n):
u, k, *vc = map(int, input().split())
G[u] = {vc[2*i]:vc[2*i+1] for i in range(k)}
d = [None] * n
heap = [[float("inf"), v] for v in range(n)]
heap[0] = [0, 0]
while heap:
heapify(heap)
d_u, u = heappop(heap)
d[u] = d_u
for i, [d_v, v] in enumerate(heap):
if v in G[u]:
d_v_new = d_u + G[u][v]
if d_v > d_v_new:
heap[i] = [d_v_new, v]
for v in range(n):
print(v, d[v])
問題は、ヒープから取り出した各節点u
に対してヒープにわたるfor
文を実行している$O(V^2)$の部分である
- ヒープにわたる
for
文を隣接リストにわたるfor
文へと狭めたいのだけれど、そうすると隣接節点v
がヒープのどこに格納されているか分からず、v
の重み$d$をうまく小さく更新できない。pythonのヒープはOrderedDictなどではなくリストを使っており、値の検索には$O(n)$かかってしまうため、隣接節点v
がヒープのどこに格納されているかを毎回検索していては計算量は結局下がらない - そこで、ヒープに含まれている節点
v
の重みの更新は諦め、重みが小さく更新された節点v
をその都度プッシュでヒープに追加していくことにする。ヒープには重みが異なる複数の同一節点が含まれることになるが、最小ヒープなので最も最近追加した最も軽い節点から取り出されることになる。削除され損ねて二回目以降に取り出された節点は、一回目に取り出されたときより重みが大きいため、if d[v] > d_v
部分がTrue
とならずにとくに何もしないので無害
以上を踏まえ、下のように書き直すとIIも通る:
from heapq import heappop, heappush
n = int(input())
G = [None] * n
for _ in range(n):
u, k, *vc = map(int, input().split())
G[u] = {vc[2*i]:vc[2*i+1] for i in range(k)}
d = [0] + [float("inf")] * (n-1)
heap = [[0, 0]]
while heap:
d_u, u = heappop(heap)
for v in G[u]:
d_v = d_u + G[u][v]
if d[v] > d_v:
d[v] = d_v
heappush(heap, [d[v], v])
for v in range(n):
print(v, d[v])
- AOJでは使えないけれど、AtCoderなどscipyが使える環境下では、自分でDijkstraのアルゴリズムのコードを書かなくても関数
scipy.sparse.csgraph.dijkstra()
を使えばよい