和を取る UnionFind

各頂点に値が振られているグラフで、連結成分内の全頂点の和を管理しながらグラフを連結させられます。

UnionFind を定義

class UnionFind:
  def __init__(self, n): pass
  def root(self, x): pass
  def unite(self, x, y): pass

root(x): x の属する UF 木の根を返します
unite(x, y): x の UF 木と y の UF 木を連結させます

UnionFind を拡張して和を管理できるようにする

class UFSum:
  def __init__(self, n):
    self.uf = UnionFind(n)
    self.val = [0] * (n + 1)
  def unite(self, x, y):
    rtx = self.uf.root(x)
    rty = self.uf.root(y)
    self.uf.unite(x, y)
    self.val[self.uf.root(x)] = self.val[rtx] + self.val[rty]
  def add(self, x, v):
    rt = self.uf.root(x)
    self.val[rt] += v
  def sum(self, x):
    return self.val[self.uf.root(x)]

前項のメソッドを実装している UnionFind データ構造であればこのデータ構造を動かすことができます。（トレイトがほしい！）

unite(x, y): x の UF 木と y の UF 木を連結させます
add(x, v): 頂点 x の値に v を足します
sum(x): x の属する連結成分の全頂点の和を返します

応用

足し算に限らず、可換性・結合性のある演算なら add(x, v), sum(x) に使うことができます。

おわり

よく使います。