初心者のプログラミング日記

プログラミング初心者の日記

プログラミングに関することを書いていきます。

クラスカル法

今回はクラスカル法をやっていきます。

使用する問題は以下になります。
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=GRL_2_A&lang=ja

まず、入力から

from heapq import heappush,heappop,heapify
V,E=map(int,input().split())
Graph=[[] for _ in range(V)]#隣接リスト
visit=[False]*(V)#頂点を訪れたか
visit[0]=True#頂点0から始めるのでそこは訪れていることにする

for i in range(E):
    s,t,w=map(int,input().split())
    Graph[s].append([w,t])#sからtに行くコスト
    Graph[t].append([w,s])#今回は無向グラフなのでtからsに行くことも可能
    
que=[i for i in Graph[0]]#キュー
heapify(que)#リストを優先度付きキューに変換
ans=0#答え

今回はheapqを使っています。
最小全域木を求めるのでコストが小さい辺を通っていくことが最善のため。
個人的にheapqでわからないことがあるのですが、

[[1, 3], [3, 2], [7, 3], [3, 4], [6, 5]] 実際はこうなる
[[1,3],[3,2],[3,4],[6,5],[7,3]] こうなると思ってた

上記のような並び替えになるんですが、調べて見たらちょっと理解できたので脱線しますが書いておきます。
簡単な例として以下の配列を優先度付きキューに変換したとします

[-1, 0, 1, 2, 3, 4, -10]

そうしたら以下の配列になります

[-10, 0, -1, 2, 3, 4, 1]

私は以下のような配列になるのかなと思ってたんですが違いました

[-10, -1, 0, 1, 2, 3, 4]

公式サイトを見てみとheapq以下の木構造になっていることがわかりました。

f:id:nasubiFX:20201103163020p:plain
公式サイトより

https://docs.python.org/ja/3/library/heapq.html

さっきの配列で木構造を作ってみると以下のようになります。
f:id:nasubiFX:20201103165421p:plain
結局上にある数値より下の数値が大きくなればいいので以下の配列になっても問題ないということだと思います。
自信はありませんが。

[-10, 0, -1, 2, 3, 4, 1]

一通り調べたのですが2次元配列の場合は乗っていませんでした。おそらく各配列の1番目の値で並び替えているとおもいます。(多分)
なので、1番目の値に各辺のコストを入れることによって、最小全域木を求めることができる。

参考記事
https://python.ms/binary-heap/#_2-%E3%82%B5%E3%83%B3%E3%83%95%E3%82%9A%E3%83%AB%E3%82%B3%E3%83%BC%E3%83%88%E3%82%99

続いて処理をかいていきます。

while que:
    c,v= heappop(que)#次に訪れる頂点とそこにいくためのコスト
    if visit[v]:#すでに訪れていたら処理を飛ばす
        continue
    visit[v]=True#訪れたらTrue
    ans+=c#答えにコストを足していく
    for e,t in Graph[v]:#頂点vに隣接する頂点にいく
        if visit[t]:#すでに訪れていたら処理を飛ばす
            continue
        heappush(que,[e,t])#訪れる場合キューに追加
        
print(ans)

最小全域木を求める場合各頂点は1回しか訪れないのでvisitで訪れたか判断します。
heapqを使っているので辺のコストが低い順に訪れるようになっている。

コード全体

from heapq import heappush,heappop,heapify
V,E=map(int,input().split())
Graph=[[] for _ in range(V)]#隣接リスト
visit=[False]*(V)#頂点を訪れたか
visit[0]=True#頂点0から始めるのでそこは訪れていることにする

for i in range(E):
    s,t,w=map(int,input().split())
    Graph[s].append([w,t])#sからtに行くコスト
    Graph[t].append([w,s])#今回は無向グラフなのでtからsに行くことも可能
    
que=[i for i in Graph[0]]#キュー
heapify(que)#リストを優先度付きキューに変換
ans=0#答え

while que:
    c,v= heappop(que)#次に訪れる頂点とそこにいくためのコスト
    if visit[v]:#すでに訪れていたら処理を飛ばす
        continue
    visit[v]=True#訪れたらTrue
    ans+=c#答えにコストを足していく
    for e,t in Graph[v]:#頂点vに隣接する頂点にいく
        if visit[t]:#すでに訪れていたら処理を飛ばす
            continue
        heappush(que,[e,t])#訪れる場合キューに追加
        
print(ans)

参考記事
https://tjkendev.github.io/procon-library/python/graph/min_st_prim.html