Goによる競プロライブラリ~Union-Find木~

内容

自身の記事である競プロのライブラリ整理~Union-Find木~のPythonコードを元に、GoによってUnion-Find木を実装しました。参考になれば幸いです。

コード

下記が実装になります。

/*
n: 要素数
parent[i]: iの親となる要素
siz[i]: iを含む素集合のサイズ
group[i]j: 代表元をiとする素集合に含まれる要素の配列j
*/
type UnionFind struct {
	n      int
	parent []int
	siz    []int
	group  map[int][]int
}

/*
コンストラクタ
任意のiについて、parent[i]=i、siz[i]=1、として初期化
group[i]jは空のmapとして初期化
*/
func NewUnionFind(n int) *UnionFind {
	uf := new(UnionFind)
	uf.n = n
	uf.parent = make([]int, n)
	uf.siz = make([]int, n)
	uf.group = make(map[int][]int)
	for i := 0; i < n; i++ {
		uf.parent[i] = i
		uf.siz[i] = 1
	}
	return uf
}

/*
要素xの属する木（素集合）の根の値を取得し、インデックスを返す
同時に経路圧縮も行う
*/
func (uf *UnionFind) root(x int) int {
	if uf.parent[x] == x {
		return x
	} else {
		uf.parent[x] = uf.root(uf.parent[x])
		return uf.parent[x]
	}
}

/*
要素xと要素yを含む素集合を併合する
*/
func (uf *UnionFind) unite(x, y int) {
	rx, ry := uf.root(x), uf.root(y)
	if rx == ry {
		return
	}
	if uf.siz[rx] < uf.siz[ry] {
		rx, ry = ry, rx
	}
	uf.siz[rx] += uf.siz[ry]
	uf.parent[ry] = rx
	return
}

/*
要素xと要素yが同じ素集合に属するかを判定し、真偽値を返す
*/
func (uf *UnionFind) same(x, y int) bool {
	rx, ry := uf.root(x), uf.root(y)
	return rx == ry
}

func (uf *UnionFind) size(x int) int {
	return uf.siz[uf.root(x)]
}

/*
現在の素集合の状態に基づいてgroupを更新する
O(uf.n)なので注意
*/
func (uf *UnionFind) grouping() {
	for i := 0; i < uf.n; i++ {
		uf.root(i)
	}
	for i := 0; i < uf.n; i++ {
		uf.group[uf.parent[i]] = append(uf.group[uf.parent[i]], i)
	}
}

/*
素集合を初期化する
*/
func (uf *UnionFind) clear() {
	for i := 0; i < uf.n; i++ {
		uf.parent[i] = i
		uf.siz = make([]int, uf.n)
	}
	uf.group = make(map[int][]int)
}