Union-find.

export type Step = {
	/** up: read one parent link. point: repoint a node straight at its root. link: put one root under another. */
	kind: 'up' | 'point' | 'link';
	node: string;
	/** The parent that was read, the root pointed at, or the root linked under. */
	to: string;
	/** For link: the size of the merged set. */
	size: number;
};
export type Options = { bySize?: boolean; compress?: boolean };

// Disjoint sets as a forest. Every item points at a parent, and a root points at itself; two
// items are in the same set when their roots match. Union by size puts the smaller tree under
// the larger, so no tree grows deeper than log2 of its size. Path compression repoints every
// item on a find straight at the root, so the next find from there is one step.
export class DisjointSets {
	readonly bySize: boolean;
	readonly compress: boolean;
	#parent = new Map<string, string>();
	#size = new Map<string, number>();
	#sets = 0;

	constructor({ bySize = true, compress = true }: Options = {}) {
		this.bySize = bySize;
		this.compress = compress;
	}

	has(id: string): boolean {
		return this.#parent.has(id);
	}

	add(id: string): void {
		if (this.#parent.has(id)) throw new RangeError(`${id} is already in a set`);
		this.#parent.set(id, id);
		this.#size.set(id, 1);
		this.#sets++;
	}

	find(id: string, steps?: Step[]): string {
		this.#need(id);
		return this.#root(id, steps);
	}

	// Returns false when the two were already in one set.
	union(a: string, b: string, steps?: Step[]): boolean {
		this.#need(a, b);
		let big = this.#root(a, steps);
		let small = this.#root(b, steps);
		if (big === small) return false;
		if (this.bySize && this.#size.get(big)! < this.#size.get(small)!) [big, small] = [small, big];
		const size = this.#size.get(big)! + this.#size.get(small)!;
		this.#parent.set(small, big);
		this.#size.set(big, size);
		this.#sets--;
		steps?.push({ kind: 'link', node: small, to: big, size });
		return true;
	}

	connected(a: string, b: string, steps?: Step[]): boolean {
		this.#need(a, b);
		return this.#root(a, steps) === this.#root(b, steps);
	}

	sizeOf(id: string): number {
		this.#need(id);
		return this.#size.get(this.#root(id))!;
	}

	parentOf(id: string): string {
		this.#need(id);
		return this.#parent.get(id)!;
	}

	/** Every set, each in the order its items were added, sets in the order of their first item. */
	groups(): string[][] {
		const at = new Map<string, string[]>();
		for (const id of this.#parent.keys()) {
			const root = this.#root(id);
			if (!at.has(root)) at.set(root, []);
			at.get(root)!.push(id);
		}
		return [...at.values()];
	}

	items(): string[] {
		return [...this.#parent.keys()];
	}

	get count(): number {
		return this.#sets;
	}

	#need(...ids: string[]): void {
		for (const id of ids) if (!this.#parent.has(id)) throw new RangeError(`unknown item ${id}`);
	}

	// Follow parents to the root, then, with compression, repoint the path at it.
	#root(id: string, steps?: Step[]): string {
		let root = id;
		for (let up = this.#parent.get(root)!; up !== root; up = this.#parent.get(root)!) {
			steps?.push({ kind: 'up', node: root, to: up, size: 0 });
			root = up;
		}
		if (this.compress)
			for (let node = id; this.#parent.get(node) !== root;) {
				const next = this.#parent.get(node)!;
				this.#parent.set(node, root);
				steps?.push({ kind: 'point', node, to: root, size: 0 });
				node = next;
			}
		return root;
	}
}

// Step is one move. Kind is "up" (read one parent link), "point" (repoint a node straight
// at its root), or "link" (put one root under another). To is the parent read, the root
// pointed at, or the root linked under; Size is the merged set's size for a link.
type Step struct {
	Kind string `json:"kind"`
	Node string `json:"node"`
	To   string `json:"to"`
	Size int    `json:"size"`
}

// DisjointSets is a forest. Every item points at a parent, and a root points at itself;
// two items are in the same set when their roots match. Union by size puts the smaller
// tree under the larger, so no tree grows deeper than log2 of its size. Path compression
// repoints every item on a find straight at the root, so the next find from there is one step.
type DisjointSets struct {
	BySize   bool
	Compress bool
	parent   map[string]string
	size     map[string]int
	order    []string
	sets     int
}

func NewDisjointSets(bySize, compress bool) *DisjointSets {
	return &DisjointSets{BySize: bySize, Compress: compress, parent: map[string]string{}, size: map[string]int{}}
}

func (d *DisjointSets) Has(id string) bool {
	_, ok := d.parent[id]
	return ok
}

func (d *DisjointSets) Add(id string) error {
	if d.Has(id) {
		return fmt.Errorf("%s is already in a set", id)
	}
	d.parent[id] = id
	d.size[id] = 1
	d.order = append(d.order, id)
	d.sets++
	return nil
}

func (d *DisjointSets) Find(id string, steps *[]Step) (string, error) {
	if err := d.need(id); err != nil {
		return "", err
	}
	return d.root(id, steps), nil
}

// Union returns false when the two were already in one set.
func (d *DisjointSets) Union(a, b string, steps *[]Step) (bool, error) {
	if err := d.need(a, b); err != nil {
		return false, err
	}
	big, small := d.root(a, steps), d.root(b, steps)
	if big == small {
		return false, nil
	}
	if d.BySize && d.size[big] < d.size[small] {
		big, small = small, big
	}
	d.parent[small] = big
	d.size[big] += d.size[small]
	d.sets--
	record(steps, "link", small, big, d.size[big])
	return true, nil
}

func (d *DisjointSets) Connected(a, b string, steps *[]Step) (bool, error) {
	if err := d.need(a, b); err != nil {
		return false, err
	}
	return d.root(a, steps) == d.root(b, steps), nil
}

func (d *DisjointSets) SizeOf(id string) (int, error) {
	if err := d.need(id); err != nil {
		return 0, err
	}
	return d.size[d.root(id, nil)], nil
}

func (d *DisjointSets) ParentOf(id string) (string, error) {
	if err := d.need(id); err != nil {
		return "", err
	}
	return d.parent[id], nil
}

// Groups returns every set, each in the order its items were added, sets in the order of
// their first item.
func (d *DisjointSets) Groups() [][]string {
	at := map[string]int{}
	groups := [][]string{}
	for _, id := range d.order {
		root := d.root(id, nil)
		i, seen := at[root]
		if !seen {
			i = len(groups)
			at[root] = i
			groups = append(groups, nil)
		}
		groups[i] = append(groups[i], id)
	}
	return groups
}

func (d *DisjointSets) Items() []string { return slices.Clone(d.order) }
func (d *DisjointSets) Count() int      { return d.sets }

func (d *DisjointSets) need(ids ...string) error {
	for _, id := range ids {
		if !d.Has(id) {
			return fmt.Errorf("unknown item %s", id)
		}
	}
	return nil
}

// root follows parents to the root, then, with compression, repoints the path at it.
func (d *DisjointSets) root(id string, steps *[]Step) string {
	root := id
	for d.parent[root] != root {
		record(steps, "up", root, d.parent[root], 0)
		root = d.parent[root]
	}
	if d.Compress {
		for node := id; d.parent[node] != root; {
			next := d.parent[node]
			d.parent[node] = root
			record(steps, "point", node, root, 0)
			node = next
		}
	}
	return root
}

func record(steps *[]Step, kind, node, to string, size int) {
	if steps != nil {
		*steps = append(*steps, Step{kind, node, to, size})
	}
}