競プロ初中級者100問

GRL2A 最小全域木

この記事で使うアルゴリズム

プリム法

プリム法のアルゴリズム解説(wikipedia)

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_2_A

この問題は最小全域木を求める問題です。プリム法を用いて解くことができます。

優先度付きキューを使って高速化しました。

コード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import heapq
V, E = map(int, input().split()) # 頂点、辺
adj = [ [] for v in range(V)] # 隣接リスト
for e in range(E):
    s, t, w = map(int, input().split())
    adj[s].append((t,w))
    adj[t].append((s,w))

visited = [0] * V # 訪問フラグ
pq = [] # 優先度付きキュー
for t, w in adj[0]: # 始点はどこでも良いので、0とする
    heapq.heappush(pq, (w, t))
visited[0] = 1 # 始点の訪問フラグを立てる

ans = 0
while(pq):
    w, t = heapq.heappop(pq)
    if visited[t]: # 訪問済みならスキップ
        continue
    visited[t] = 1 # 訪問フラグを立てる
    ans += w # スコアに加算
    for s, w in adj[t]: # 隣接する頂点を列挙
        if visited[s]==0: # 未訪問なら探索候補としてpqに追加
            heapq.heappush(pq, (w, s))
print (ans)

記事情報

  • 投稿日：2020年4月28日
  • 最終更新日：2020年5月10日