gno.land/p/moul/x/daily/disjointset/v0⧉

1// Package disjointset is union-find (a disjoint-set forest) as a pure,2// reusable package: it tracks a partition of [0, n) into disjoint groups and3// answers "are these two in the same group?" in near-constant time.4//5// Both classic optimisations are implemented, and they matter together: path6// compression flattens a tree on every Find, union by rank keeps the shallower7// tree under the deeper one. With both, operations are O(α(n)) — inverse8// Ackermann, effectively constant. With neither, a chain of unions degrades to9// O(n) per query, which on chain is the difference between a cheap call and an10// out-of-gas one.11//12// A live demo of this package is at13// [r/moul/x/daily/disjointsetdemo](/r/moul/x/daily/disjointsetdemo/v0).14package disjointset1516// MaxN bounds a set so allocation stays predictable.17const MaxN = 1 << 161819// DisjointSet is a partition of [0, n) into disjoint groups.20type DisjointSet struct {21	parent []int22	rank   []int23	groups int24}2526// New returns n singleton groups. n is clamped to [0, MaxN].27func New(n int) *DisjointSet {28	if n < 0 {29		n = 030	}31	if n > MaxN {32		n = MaxN33	}34	d := &DisjointSet{parent: make([]int, n), rank: make([]int, n), groups: n}35	for i := 0; i < n; i++ {36		d.parent[i] = i // every element starts as its own root37	}38	return d39}4041// Len returns the number of elements.42func (d *DisjointSet) Len() int { return len(d.parent) }4344// Groups returns how many disjoint groups remain.45func (d *DisjointSet) Groups() int { return d.groups }4647// InRange reports whether i is a valid element.48func (d *DisjointSet) InRange(i int) bool { return i >= 0 && i < len(d.parent) }4950// Find returns the representative of i's group, or -1 when i is out of range.51//52// Path compression: every node visited is re-pointed straight at the root, so53// the next Find on any of them is O(1). Done iteratively rather than54// recursively — a deep chain would otherwise risk the call stack.55func (d *DisjointSet) Find(i int) int {56	if !d.InRange(i) {57		return -158	}59	root := i60	for d.parent[root] != root {61		root = d.parent[root]62	}63	for d.parent[i] != root { // second pass: re-point everything at the root64		next := d.parent[i]65		d.parent[i] = root66		i = next67	}68	return root69}7071// Union merges the groups of a and b and reports whether they were merged.72// False means they were already together, or an index was out of range.73func (d *DisjointSet) Union(a, b int) bool {74	ra, rb := d.Find(a), d.Find(b)75	if ra < 0 || rb < 0 || ra == rb {76		return false77	}78	// Union by rank: hang the shallower tree off the deeper one so depth only79	// grows when both sides are equally deep.80	if d.rank[ra] < d.rank[rb] {81		ra, rb = rb, ra82	}83	d.parent[rb] = ra84	if d.rank[ra] == d.rank[rb] {85		d.rank[ra]++86	}87	d.groups--88	return true89}9091// Connected reports whether a and b are in the same group. Out-of-range92// indices are not connected to anything, including themselves.93func (d *DisjointSet) Connected(a, b int) bool {94	ra, rb := d.Find(a), d.Find(b)95	return ra >= 0 && ra == rb96}9798// Size returns how many elements share i's group, or 0 when out of range.99func (d *DisjointSet) Size(i int) int {100	r := d.Find(i)101	if r < 0 {102		return 0103	}104	n := 0105	for j := 0; j < len(d.parent); j++ {106		if d.Find(j) == r {107			n++108		}109	}110	return n111}112113// Partition returns the groups, each sorted ascending, ordered by their114// smallest member — deterministic regardless of the union order, which is what115// makes it safe to render.116func (d *DisjointSet) Partition() [][]int {117	byRoot := map[int][]int{}118	order := []int{}119	for i := 0; i < len(d.parent); i++ {120		r := d.Find(i)121		if _, seen := byRoot[r]; !seen {122			order = append(order, r) // first sighting is the smallest member123		}124		byRoot[r] = append(byRoot[r], i)125	}126	out := [][]int{}127	for _, r := range order { // iterate the slice, never the map128		out = append(out, byRoot[r])129	}130	return out131}132