Kruskal’s minimum spanning tree.

// Kruskal's minimum spanning tree: connect every spot on a festival site with the least cable.
// Every cable run is two-way and measured in whole meters.
export const MAX_SPOTS = 64;
export const MAX_RUNS = 512;
export const MAX_METRES = 1000;

export type Run = { a: string; b: string; metres: number };
export type Site = { spots: string[]; runs: Run[] };
// The runs to lay, their total length, and the groups of spots they join. One group means every
// spot is connected; more than one means the runs ran out first.
export type Network = { runs: Run[]; metres: number; groups: string[][] };

export type CableErrorCode = 'too-big' | 'bad-spot' | 'unknown-spot' | 'bad-metres';

export class CableError extends Error {
	readonly code: CableErrorCode;
	constructor(code: CableErrorCode, message: string) {
		super(message);
		this.name = 'CableError';
		this.code = code;
	}
}

// Union-find over spot numbers, with union by size and path compression.
export class Groups {
	#parent: number[];
	#size: number[];
	count: number;

	constructor(spots: number) {
		this.#parent = Array.from({ length: spots }, (_, i) => i);
		this.#size = Array.from({ length: spots }, () => 1);
		this.count = spots;
	}

	find(spot: number): number {
		let root = spot;
		while (this.#parent[root] !== root) root = this.#parent[root];
		while (this.#parent[spot] !== root) {
			const next = this.#parent[spot];
			this.#parent[spot] = root;
			spot = next;
		}
		return root;
	}

	// Returns false when the two spots are already in one group.
	union(a: number, b: number): boolean {
		let big = this.find(a);
		let small = this.find(b);
		if (big === small) return false;
		if (this.#size[big] < this.#size[small]) [big, small] = [small, big];
		this.#parent[small] = big;
		this.#size[big] += this.#size[small];
		this.count--;
		return true;
	}

	// Every group, spots in site order, groups in the order of their first spot.
	list(spots: string[]): string[][] {
		const at = new Map<number, string[]>();
		spots.forEach((spot, i) => {
			const root = this.find(i);
			if (!at.has(root)) at.set(root, []);
			at.get(root)!.push(spot);
		});
		return [...at.values()];
	}
}

// Check the site and number its spots in the order they are listed.
export function numberSpots(site: Site, laid: Run[] = []): Map<string, number> {
	if (site.spots.length > MAX_SPOTS || site.runs.length + laid.length > MAX_RUNS)
		throw new CableError('too-big', `up to ${MAX_SPOTS} spots and ${MAX_RUNS} runs`);
	const index = new Map<string, number>();
	for (const spot of site.spots) {
		if (!/^[a-z][a-z0-9-]{0,23}$/.test(spot) || index.has(spot))
			throw new CableError('bad-spot', `spot ids are unique lowercase slugs: ${spot}`);
		index.set(spot, index.size);
	}
	for (const { a, b, metres } of [...site.runs, ...laid]) {
		if (!index.has(a) || !index.has(b))
			throw new CableError('unknown-spot', `a run joins spots that don't exist: ${a}, ${b}`);
		if (!Number.isInteger(metres) || metres < 1 || metres > MAX_METRES)
			throw new CableError('bad-metres', `metres are whole numbers from 1 to ${MAX_METRES}`);
	}
	return index;
}

// Shortest first. Equal lengths keep their listed order, so both languages pick the same runs.
export const byLength = (runs: Run[]): Run[] => [...runs].sort((x, y) => x.metres - y.metres);

export function cheapestNetwork(site: Site): Network {
	const index = numberSpots(site);
	const groups = new Groups(site.spots.length);
	const runs: Run[] = [];
	let metres = 0;
	for (const run of byLength(site.runs)) {
		if (runs.length === site.spots.length - 1) break; // n spots are joined by n − 1 runs
		// Both ends already in one group: this run would close a loop, and it is the longest run
		// on that loop, because every other run on it was taken earlier.
		if (!groups.union(index.get(run.a)!, index.get(run.b)!)) continue;
		runs.push(run);
		metres += run.metres;
	}
	return { runs, metres, groups: groups.list(site.spots) };
}

// Kruskal's minimum spanning tree: connect every spot on a festival site with the least cable.
// Every cable run is two-way and measured in whole meters.
const (
	MaxSpots  = 64
	MaxRuns   = 512
	MaxMetres = 1000
)

type Run struct {
	A, B   string
	Metres int
}

type Site struct {
	Spots []string
	Runs  []Run
}

// Network is the runs to lay, their total length, and the groups of spots they join. One group
// means every spot is connected; more than one means the runs ran out first.
type Network struct {
	Runs   []Run
	Metres int
	Groups [][]string
}

type CableError struct{ Code, Message string }

func (e *CableError) Error() string { return e.Message }

// groups is union-find over spot numbers, with union by size and path compression.
type groups struct {
	parent, size []int
	count        int
}

func newGroups(spots int) *groups {
	g := &groups{parent: make([]int, spots), size: make([]int, spots), count: spots}
	for i := range g.parent {
		g.parent[i] = i
		g.size[i] = 1
	}
	return g
}

func (g *groups) find(spot int) int {
	root := spot
	for g.parent[root] != root {
		root = g.parent[root]
	}
	for g.parent[spot] != root {
		next := g.parent[spot]
		g.parent[spot] = root
		spot = next
	}
	return root
}

// union returns false when the two spots are already in one group.
func (g *groups) union(a, b int) bool {
	big, small := g.find(a), g.find(b)
	if big == small {
		return false
	}
	if g.size[big] < g.size[small] {
		big, small = small, big
	}
	g.parent[small] = big
	g.size[big] += g.size[small]
	g.count--
	return true
}

// list returns every group, spots in site order, groups in the order of their first spot.
func (g *groups) list(spots []string) [][]string {
	out := [][]string{}
	at := map[int]int{}
	for i, spot := range spots {
		root := g.find(i)
		k, seen := at[root]
		if !seen {
			k = len(out)
			at[root] = k
			out = append(out, nil)
		}
		out[k] = append(out[k], spot)
	}
	return out
}

var spotID = regexp.MustCompile(`^[a-z][a-z0-9-]{0,23}$`)

// NumberSpots checks the site and numbers its spots in the order they are listed.
func NumberSpots(site Site, laid []Run) (map[string]int, error) {
	if len(site.Spots) > MaxSpots || len(site.Runs)+len(laid) > MaxRuns {
		return nil, &CableError{"too-big", fmt.Sprintf("up to %d spots and %d runs", MaxSpots, MaxRuns)}
	}
	index := make(map[string]int, len(site.Spots))
	for _, spot := range site.Spots {
		if _, seen := index[spot]; seen || !spotID.MatchString(spot) {
			return nil, &CableError{"bad-spot", "spot ids are unique lowercase slugs: " + spot}
		}
		index[spot] = len(index)
	}
	for _, run := range slices.Concat(site.Runs, laid) {
		_, knownA := index[run.A]
		_, knownB := index[run.B]
		if !knownA || !knownB {
			return nil, &CableError{"unknown-spot", fmt.Sprintf("a run joins spots that don't exist: %s, %s", run.A, run.B)}
		}
		if run.Metres < 1 || run.Metres > MaxMetres {
			return nil, &CableError{"bad-metres", fmt.Sprintf("metres are whole numbers from 1 to %d", MaxMetres)}
		}
	}
	return index, nil
}

// byLength sorts shortest first. Equal lengths keep their listed order, so both languages pick
// the same runs.
func byLength(runs []Run) []Run {
	sorted := slices.Clone(runs)
	slices.SortStableFunc(sorted, func(x, y Run) int { return x.Metres - y.Metres })
	return sorted
}

func CheapestNetwork(site Site) (Network, error) {
	index, err := NumberSpots(site, nil)
	if err != nil {
		return Network{}, err
	}
	g := newGroups(len(site.Spots))
	network := Network{Runs: []Run{}}
	for _, run := range byLength(site.Runs) {
		if len(network.Runs) == len(site.Spots)-1 {
			break // n spots are joined by n − 1 runs
		}
		// Both ends already in one group: this run would close a loop, and it is the longest run
		// on that loop, because every other run on it was taken earlier.
		if !g.union(index[run.A], index[run.B]) {
			continue
		}
		network.Runs = append(network.Runs, run)
		network.Metres += run.Metres
	}
	network.Groups = g.list(site.Spots)
	return network, nil
}