prim's algorithm

date: 2021-04-25 excerpt: プリム法について

tag: algorithm network graph prim

プリム(prim)法について

最小全域木探索アルゴリズムの一つ
- ダイクストラ法の特殊系とみなせる
ある支点から最もコストが安い経路を一つ選ぶ
- この経路を追加した上で最もその中から最もコストが安いものを選ぶ
  - heapで優先順位を管理しながら探索すると現実的な時間で解くことができる

例; コストを反転した最少全域木問題

問題

いろはちゃんコンテスト Day2; D - 楽しすぎる家庭菜園

解説

最大の流れを持つ川を残し、他をカットする = 川の流れを反転して負のコストにする

解答

import heapq
import collections
N, M = map(int, input().split())

G = collections.defaultdict(list)
for node_num in range(M):
    u, v, c = map(int, input().split())
    u-=1; v-=1
    G[u].append((v, -c, node_num))
    G[v].append((u, -c, node_num))
que = []
# 初期化
[heapq.heappush(que, (c, j, node_num)) for j, c, node_num in G[0]]
visited = set([0])

traces = []
while que:
    c, i, node_num = heapq.heappop(que)
    if i in visited:
        continue
    visited.add(i)
    traces.append(node_num+1)
    for (j, nc, nnode_num) in G[i]:
        if j in visited:
            continue
        heapq.heappush(que, (nc, j, nnode_num))
traces.sort()
print(*traces, sep="\n")

例; 最小全域木問題

問題

典型アルゴリズム問題集; F - 最小全域木問題

解答

import heapq
import collections
N, M = map(int, input().split())
G=collections.defaultdict(list)
for _ in range(M):
    u, v, c = map(int, input().split())
    G[u].append((v, c))
    G[v].append((u, c))

# 初期化
que, visited = [], set([0])
[heapq.heappush(que, (c, j)) for j, c in G[0]]

ans = 0
while que:
    c, i = heapq.heappop(que)
    if visited.__contains__(i):
        continue
    visited.add(i)
    ans += c
    for (j, nc) in G[i]:
        if visited.__contains__(j):
            continue
        heapq.heappush(que, (nc, j))
print(ans)

例; 最小コストを探索するとき

問題

AtCoder Beginner Contest 065

解説

初期化のために最初にqueを作成する段階でmarkedのような到達したノードを管理する変数を連動させることを忘れがちである
純粋に文字列をキーとして探索することもできるが、pypyでは許容できそうにない遅さになる(pypyは文字列比較等が遅い)

解答

import heapq
import collections
N = int(input())
G = collections.defaultdict(list)

XY = []
for _ in range(N):
    x,y = map(int, input().split())
    XY.append( (x,y) )


XY.sort()
for i in range(N-1):
    x, y = XY[i]
    nx, ny = XY[i+1]
    cost = min(abs(x-nx), abs(y-ny))
    G[(x,y)].append( ((nx,ny), cost) )
    G[(nx,ny)].append( ((x,y), cost) )
XY.sort(key=lambda x:x[1])
for i in range(N-1):
    x, y = XY[i]
    nx, ny = XY[i+1]
    cost = min(abs(x-nx), abs(y-ny))
    G[(x,y)].append( ((nx,ny), cost) )
    G[(nx,ny)].append( ((x,y), cost) )

que = []
start = XY[0]
# 初期化(prim法的には任意のどこからでも初めて良い)
[heapq.heappush(que, (c, xy)) for xy, c in G[start]]
visited = set([start])

acc = 0
while que:
    c, i = heapq.heappop(que)
    if i in visited:
        continue
    visited.add(i)
    acc += c
    for (j, nc) in G[i]:
        if j in visited:
            continue
        heapq.heappush(que, (nc, j))
print(acc)