LoginSignup
2
2

More than 3 years have passed since last update.

Pythonによる螺旋本・AOJの解答(ALDS1 #7~#12)

Last updated at Posted at 2020-09-11

この記事では、螺旋本・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}")
  • paidへの参照と子達chsidへの参照を保持するクラス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 BAが真なら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

以下のようにextractinsertを自分で書くと、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種類

    • 白(未発見):dfNone
    • 灰(発見済だけど未完了):dに発見時刻が入っているけれど、fNone
    • 黒(完了済):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)$で分かる。未訪問の節点はrepsNone。代表元としては、その連結成分で最初に訪れた節点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()を使えばよい
2
2
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
2
2