最小全域木

クラスカル法の前に理解しておかなければならない概念。全域木はM個の辺の中からいくつかを選んで作った、全ての頂点が連結されている木。その中でも辺の重さの合計が最小になるものを最小全域木と呼ぶ*1。

使い所の例としてはいくつかの都市をつなぐ鉄道を建設したいが年度ごとの予算に限りがある場合を想定すると良い。この場合、どこかの都市を基点としてまずは安く建設できる都市と繋いでいく。そういう場合に使えるアルゴリズムである。

クラスカル法

クラスカル法は最小全域木を求めるためのアルゴリズムだ。やりかたは簡単で、辺の重みをソートしてUnion Findしていくだけ。計算量はO(M log M + N)。Union Findは別途Wikiを参照。

では例題で確認する。

例題 ABC218 E

E - Destruction

例えば上図のような全域木があるとする。辺を取り除いた後も全域木である必要があるため、取り除ける辺は赤線の2つのいずれかとなる。この場合は500を取り除くことが答えになる。しかし「取り除く」辺を探索して実際に取り除く実装はかなり骨が折れそうだ。

そこで、クラスカル法により最小全域木を作ることを考える。この時辺の重みが小さいものから選んでUnion Findをする。その結果、選ばれなかった辺を足したものが答えとなる。

以下のように実装した。

class UnionFind:
    def __init__(self, n):
        self.n = n
        self.parent_size = [-1] * n

    """
  リーダーが同じなら何もしない
  リーダーが違う場合、グループの人数が少ない方を、大きい方のグループへ統合する
  """

    def merge(self, a, b):
        x, y = self.leader(a), self.leader(b)
        if x == y:
            return
        if abs(self.parent_size[x]) < abs(self.parent_size[y]):
            x, y = y, x
        self.parent_size[x] += self.parent_size[y]
        self.parent_size[y] = x
        return

    def same(self, a, b):
        return self.leader(a) == self.leader(b)

    """
  再帰で根にいるリーダーを探しに行く
  自分自身がリーダーなら自分を返す
  """

    def leader(self, a):
        if self.parent_size[a] < 0:
            return a
        self.parent_size[a] = self.leader(self.parent_size[a])
        return self.parent_size[a]

    """
  aが所属するグループのサイズ(人数)を返す
  """

    def size(self, a):
        return abs(self.parent_size[self.leader(a)])

    def groups(self):
        result = [[] for _ in range(self.n)]
        for i in range(self.n):
            result[self.leader(i)].append(i)
        return [r for r in result if r != []]

n,m=map(int, input().split())
edges=[]
for i in range(m):
  a,b,c=map(int, input().split())
  edges.append((c, a-1, b-1))

edges.sort()

uf=UnionFind(n)
ans=0
for c,a,b in edges:
  if not uf.same(a,b):
    uf.merge(a,b)
  elif c > 0:
    ans+=c

print(ans)

余談だが、この問題の考え方は少し納得がいっていない。というのは報酬の最大値の解釈に余地があると考えているからだ。問題文中では報酬と罰金が区別されているため、c>0の場合にだけ関心を持っていれば正解ということらしい。が、厳密には報酬-罰金が真に最大の報酬となるのではないかと思った。