塩見周子の徒然日記

自分のことを塩見周子と思い込んでいるオタクです

4/20 AtCoder(python3、プリム法)

ARC076D
プリム法とは
>>>頂点をすべて結ぶのに必要な最小のコストを、辺E,頂点VとしてO(ElogV)で求めるアルゴリズム

構造は
①「『既に訪れた頂点(ア)』と繋がる『まだ訪れていない頂点(イ)』」とを結ぶ辺のリストの中で(ここではあくまで辺が主役)、最もコストが小さいものを選ぶ
②その辺のコストを足す(辺を使うことを確定させる)
③頂点(イ)と辺でつながる頂点のうち、その頂点がまだ訪れていないのであれば、辺のリストの中に追加する
④①に戻り、リストが空になるまで続ける

というもの

import heapq
def prim_heap():
    used = [True]*n #True:不使用
    edgelist = []
    for e in edge[0]: #連結なのでedge[0]は必ず存在する
        heapq.heappush(edgelist,e)
    used[0] = False #番号0の頂点を始点とする
    res = 0 #答え
    while len(edgelist) != 0:
        minedge = heapq.heappop(edgelist)
        if not used[minedge[1]]: #取り出した辺の行き先が使用済みなら
            continue #捨てる
        v = minedge[1] #そうでないなら行き先を使用済みにする
        used[v] = False
        for e in edge[v]: #行き先の頂点がつながる辺において
            if used[e[1]]: #その先の頂点が未使用なら
                heapq.heappush(edgelist,e) #次につなげる辺の候補にする
        res += minedge[0] #コストを足す
    return res

n,w = map(int,input().split()) #n:頂点の数、w:辺の数
edge = [[] for i in range(n)] #edge[i]には、[コスト,行先]というリストがさらに内包される
for i in range(w):
    x,y,z = map(int,input().split())
    edge[x].append([z,y])
    edge[y].append([z,x])
print(prim_heap())


この問題では、全部の辺をつなげようとすると、N**2通りの辺が存在するので、N = 10**5だと間に合わない。少し考えると、かけるべき辺は、ある頂点とx座標、y座標のそれぞれで最も近い頂点との間のみでよいことが分かる(x座標の小さい方から並べた頂点ア、イ、ウに対し、ア、イ、ウすべてを訪れることを考えると、アーウとつなげる必要はなく、アーイ、イーウとつなげればよいから)
だからx,yのそれぞれでソートし、隣り合う頂点のみに辺を掛ければ、辺の数はO(N)になって間に合う

実装が重いので後でやります.......。


juppy.hatenablog.com