[Python] UnionFind ABC214D

ABC214D

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

各 $i , (1 \leq i \leq N -1)$ ついて、$f(u, v) = w_i$ となる $(u, v)$ の総数を数えることを考えます。$f(u, v) = w_i$となるための条件は、頂点 $(u, v)$ を結ぶパスに辺 $i$ が含まれ、かつそのパスに含まれる他の辺の重みが $w_i$ より小さいことです。

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

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

サンプルコード

class UF:
  def __init__(self, N):
    # vertex → root, root → size
    self.uf = [-1] * N
    self.n = N

  # find root
  def find(self, x):
    if self.uf[x] < 0: return x
    self.uf[x] = self.find(self.uf[x])
    return self.uf[x]

  def size(self, x):
    return -self.uf[self.find(x)]

  def union(self, x, y):
    x, y = self.find(x), self.find(y)
    if x == y: return
    if self.size(x) > self.size(y): x, y = y, x
    self.uf[y] += self.uf[x]
    self.uf[x] = y
    self.n -= 1

N = int(input())
E = sorted([tuple(map(int, input().split())) for _ in range(N-1)], key=lambda x: x[2])

uf = UF(N)
ans = 0
for u, v, w in E:
  u -= 1
  v -= 1
  ans += uf.size(u) * uf.size(v) * w
  uf.union(u, v)

print(ans)