disjointset.gno

  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 and
  3// answers "are these two in the same group?" in near-constant time.
  4//
  5// Both classic optimisations are implemented, and they matter together: path
  6// compression flattens a tree on every Find, union by rank keeps the shallower
  7// tree under the deeper one. With both, operations are O(α(n)) — inverse
  8// Ackermann, effectively constant. With neither, a chain of unions degrades to
  9// O(n) per query, which on chain is the difference between a cheap call and an
 10// out-of-gas one.
 11//
 12// A live demo of this package is at
 13// [r/moul/x/daily/disjointsetdemo](/r/moul/x/daily/disjointsetdemo/v0).
 14package disjointset
 15
 16// MaxN bounds a set so allocation stays predictable.
 17const MaxN = 1 << 16
 18
 19// DisjointSet is a partition of [0, n) into disjoint groups.
 20type DisjointSet struct {
 21	parent []int
 22	rank   []int
 23	groups int
 24}
 25
 26// New returns n singleton groups. n is clamped to [0, MaxN].
 27func New(n int) *DisjointSet {
 28	if n < 0 {
 29		n = 0
 30	}
 31	if n > MaxN {
 32		n = MaxN
 33	}
 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 root
 37	}
 38	return d
 39}
 40
 41// Len returns the number of elements.
 42func (d *DisjointSet) Len() int { return len(d.parent) }
 43
 44// Groups returns how many disjoint groups remain.
 45func (d *DisjointSet) Groups() int { return d.groups }
 46
 47// InRange reports whether i is a valid element.
 48func (d *DisjointSet) InRange(i int) bool { return i >= 0 && i < len(d.parent) }
 49
 50// 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, so
 53// the next Find on any of them is O(1). Done iteratively rather than
 54// recursively — a deep chain would otherwise risk the call stack.
 55func (d *DisjointSet) Find(i int) int {
 56	if !d.InRange(i) {
 57		return -1
 58	}
 59	root := i
 60	for d.parent[root] != root {
 61		root = d.parent[root]
 62	}
 63	for d.parent[i] != root { // second pass: re-point everything at the root
 64		next := d.parent[i]
 65		d.parent[i] = root
 66		i = next
 67	}
 68	return root
 69}
 70
 71// 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 false
 77	}
 78	// Union by rank: hang the shallower tree off the deeper one so depth only
 79	// grows when both sides are equally deep.
 80	if d.rank[ra] < d.rank[rb] {
 81		ra, rb = rb, ra
 82	}
 83	d.parent[rb] = ra
 84	if d.rank[ra] == d.rank[rb] {
 85		d.rank[ra]++
 86	}
 87	d.groups--
 88	return true
 89}
 90
 91// Connected reports whether a and b are in the same group. Out-of-range
 92// 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 == rb
 96}
 97
 98// 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 0
103	}
104	n := 0
105	for j := 0; j < len(d.parent); j++ {
106		if d.Find(j) == r {
107			n++
108		}
109	}
110	return n
111}
112
113// Partition returns the groups, each sorted ascending, ordered by their
114// smallest member — deterministic regardless of the union order, which is what
115// 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 member
123		}
124		byRoot[r] = append(byRoot[r], i)
125	}
126	out := [][]int{}
127	for _, r := range order { // iterate the slice, never the map
128		out = append(out, byRoot[r])
129	}
130	return out
131}