[Python] UnionFindって何！？ -AtCoder学習日記 #2

概要

今回はABC214-Dです。やっとD問題に手が出るようになりました。
問題文的にこの問題はグラフ理論の問題のように見えます。

解説文

今回詳しく書く部分(私がわからなかった部分)を赤くしています。

以下では、辺の重みは相異なると仮定します。制約にそのような条件はありませんが、添字の大小による tie-break などを行えばよいです。

各について、となるの総数を数えることを考えます。となるための条件は、頂点を結ぶパスに辺が含まれ、かつそのパスに含まれる他の辺の重みがより小さいことです。

重みがより小さい辺のみを張ったグラフを考えるとこれは森(グラフ理論の用語)となっており、とは必ず連結成分に属します。これら二つの連結成分の一方に、もう一方にが属することが、となるための条件です。

よって、辺の重みの小さい順にグラフに追加していき、Union-Findなどのデータ構造を用いて連結成分の大きさを取得することにより、この問題を得ことができます。

解説の解説

解説を読む限り、グラフ理論の用語をしっかり知っておく必要がありそうです。。。

森

我らがWikipedia先生によると

閉路を持たない (連結であるとは限らない) グラフを森 (もり、英: forest) という。

とのことです。

• 閉路は"始点と終点が同じ道"、つまりループ部分のことです。

• 連結であるは"ある2点間を結ぶ道が存在すること"です。

つまり、バラバラなところがあってもいいけど、ループ部分はないグラフですかね。
木は「連結な森」とのことです。

ここまではわかります。

連結成分

次は、連結成分についてです。

極大で連結な部分グラフは、連結成分 (connected component) という。

極大で連結な部分グラフ？？？
よくわからないので、少しずつ整理します。

連結は先ほどにも出てきた通り”ある2点間を結ぶ道が存在すること”で、部分グラフは”グラフの一部分のグラフ”です。

では、極大で連結とはどういうことなのでしょうか。

Google先生に聞いてみるとここで以下のような説明がありました。

G'がグラフGの極大で連結な部分グラフであるとは、G'を部分グラフに持つGの連結な部分グラフがG'以外にないことを言います。

つまり極大で連結な部分グラフもとい連結成分はバラバラになってない一塊のグラフってことですかね！

グラフの中に、連結でない部分があることをしっかり意識しないと理解するのも難しいです。。。

Union-Find

DSU( Disjoint Set Union, 素集合データ構造)とも言うみたいです。
ここに以下のように説明があります。とてもわかりやすいです。

Union-Find木とは，グループ分けを木構造で管理するデータ構造のこと．同じグループに属する＝同じ木に属するという木構造でグループ分けをする際，以下の二点を高速で行うことができるのがメリット．

• 要素xと要素yが同じグループに属するかどうかを判定したい
  →要素xの根と要素yの根が同じならば同じグループ，要素xの根と要素yの根が同じでないならば異なるグループにあることが分かる．

• 要素xと要素yが同じグループに属する場合，要素xの属するグループと要素yの属するグループの併合する．

このスライドが非常に分かりやすい．

そのスライドも本当にわかりやすいです。
ある集合内の要素を特定の条件に従って分けるときに効果を発揮するアルゴリズムってことみたいですね。
同じ条件のグループを同じ木でまとめて、要素同じグループかどうかは木の根が同じかどうかで判断します。
同じ条件かどうかを木の根のみで判断するため、結局どこの木の根にくっついているかのみに着目すれば良いことになります。
そのため、実際には木構造を構成するのではなく、リストで木の根がどの要素なのかのみを管理すれば良いです。
これらをPythonで実装するとこんな感じですかね。

class UnionFind:
  def __init__(self, n):
    self.n = n
    self.root_size = [-1] * n 
      #木の根が誰かを格納するリスト(自分が木の根の場合は負の数)
      #グループのサイズを木の根のところに負の数で格納する
      #  →木の根の部分には、自分が木の根であることとグループのサイズ
      #   の2つの情報が入ることになる．

  def root(self, u): #木の根が誰かを返すメソッド
    if self.root_size[u] < 0:
      return u
    else:
      self.root_size[u] = self.root(self.root_size[u])
    return self.root_size[u]

  def merge(self, u, v): #二つの要素をまとめるメソッド
    u_root = self.root(u) #木の根が誰かを求めて
    v_root = self.root(v)
    if u_root == v_root: #一致していれば、すでにまとめられている
      return True
    else: #一致していなければ、
      self.root_size[v_root] += self.root_size[u_root] #片方のサイズをもう片方に足して、
      self.root_size[u_root] = v_root #木の根を変更する
      return True

  def size(self, u): #ある要素のいるグループのサイズを返すメソッド
    #木の根が誰かを求めて、そのサイズをリストから求める．
    return -self.root_size[self.root(u)]

このように実装できるとおもいます。

ACになる解答

UnionFindのPythonでの実装はもう書いたのでそれ以外の部分です。

n = int(input())
edge = []

for _ in range(n-1): #一旦辺の情報を保存して
  u, v, w = map(int, input().split())
  edge.append([u-1, v-1, w])

edge = sorted(edge, key=lambda x: x[2])
#重みが小さい順に分ける
uf = UnionFind(n)
ans = 0

for i in range(n-1):
  [u, v, w] = edge[i]
  ans += w * uf.size(u) * uf.size(v)
  #UnionFind木のそれぞれの要素が含まれるグループのサイズをかける
  # →f(u, v)がその重みになる個数がわかる。
  # →uのグループに含まれる頂点からvのグループに含まれる頂点までを結ぶ道で
  #  u-vの辺を通る場合の数
  # *重みが少ない順に処理しているため、常に全てのグループは
  #  処理する時の重みよりも小さい重みの辺のみを持つ
  uf.merge(u, v)
  #uの入ったグループと, vの入ったグループを同じグループにまとめる。
  
print(ans)

と、このような感じです。これでしっかりとACが出るはずです。

重みが小さい順に処理することで、ある重みよりも小さい辺だけで繋がった連結成分だけを考えられるのは、なかなか出てこない発想な気がします。
いや．．．ただ勉強不足なだけか．．．？

まとめ

• UnionFindは要素をグループ分けしたいときに有効
• グラフ理論はちゃんと勉強したい
• AtCoderが解説している便利なページも存在する