Minimum Spanning Trees

def kruskal(n: int, edges: list[list[int]]) -> list[list[int]]:
    """Minimum spanning tree of the nodes 0..n-1, as a list of [u, v, weight]
    edges. A disconnected graph gets a spanning forest with fewer than n - 1 edges."""
    parent = list(range(n))  # A small disjoint set: every node starts alone.
    mst = []
    # Cheapest edges first. The sort is stable, so equal weights keep their input order.
    for u, v, weight in sorted(edges, key=lambda edge: edge[2]):
        root_u, root_v = find(parent, u), find(parent, v)
        if root_u == root_v:
            continue  # Already connected: this edge would close a cycle.
        parent[root_u] = root_v  # Merge the two components.
        mst.append([u, v, weight])
        if len(mst) == n - 1:
            break  # A spanning tree has exactly n - 1 edges.
    return mst


def find(parent: list[int], x: int) -> int:
    """Root of x's set. Path halving: point each node at its grandparent on the way up."""
    while parent[x] != x:
        parent[x] = parent[parent[x]]
        x = parent[x]
    return x


if __name__ == "__main__":
    edges = [[2, 3, 5], [0, 1, 1], [3, 4, 6], [0, 3, 3], [1, 3, 2], [0, 2, 4]]
    mst = kruskal(5, edges)
    print(mst)  # [[0, 1, 1], [1, 3, 2], [0, 2, 4], [3, 4, 6]]
    print(sum(weight for _, _, weight in mst))  # 13
    # Two separate pieces: only 2 edges for 4 nodes, so there's no spanning tree.
    print(kruskal(4, [[0, 1, 3], [2, 3, 2]]))  # [[2, 3, 2], [0, 1, 3]]

// Minimum spanning tree of the nodes 0..n-1, as a list of [u, v, weight]
// edges. A disconnected graph gets a spanning forest with fewer than n - 1 edges.
export function kruskal(n: number, edges: number[][]): number[][] {
  const parent = Array.from({ length: n }, (_, i) => i); // A small disjoint set: every node starts alone.
  const mst: number[][] = [];
  // Cheapest edges first. The sort is stable, so equal weights keep their input order.
  const sorted = [...edges].sort((a, b) => a[2] - b[2]);
  for (const [u, v, weight] of sorted) {
    const rootU = find(parent, u);
    const rootV = find(parent, v);
    if (rootU === rootV) continue; // Already connected: this edge would close a cycle.
    parent[rootU] = rootV; // Merge the two components.
    mst.push([u, v, weight]);
    if (mst.length === n - 1) break; // A spanning tree has exactly n - 1 edges.
  }
  return mst;
}

// Root of x's set. Path halving: point each node at its grandparent on the way up.
function find(parent: number[], x: number): number {
  while (parent[x] !== x) {
    parent[x] = parent[parent[x]];
    x = parent[x];
  }
  return x;
}

const edges = [[2, 3, 5], [0, 1, 1], [3, 4, 6], [0, 3, 3], [1, 3, 2], [0, 2, 4]];
const mst = kruskal(5, edges);
console.log(JSON.stringify(mst)); // [[0,1,1],[1,3,2],[0,2,4],[3,4,6]]
console.log(mst.reduce((total, edge) => total + edge[2], 0)); // 13
// Two separate pieces: only 2 edges for 4 nodes, so there's no spanning tree.
console.log(JSON.stringify(kruskal(4, [[0, 1, 3], [2, 3, 2]]))); // [[2,3,2],[0,1,3]]

package main

import (
	"cmp"
	"fmt"
	"slices"
)

// kruskal returns the minimum spanning tree of the nodes 0..n-1 as a list of
// {u, v, weight} edges. A disconnected graph gets a spanning forest with
// fewer than n - 1 edges.
func kruskal(n int, edges [][3]int) [][3]int {
	parent := make([]int, n) // A small disjoint set: every node starts alone.
	for i := range parent {
		parent[i] = i
	}
	mst := [][3]int{}
	// Cheapest edges first. The sort is stable, so equal weights keep their input order.
	sorted := slices.Clone(edges)
	slices.SortStableFunc(sorted, func(a, b [3]int) int { return cmp.Compare(a[2], b[2]) })
	for _, edge := range sorted {
		rootU, rootV := find(parent, edge[0]), find(parent, edge[1])
		if rootU == rootV {
			continue // Already connected: this edge would close a cycle.
		}
		parent[rootU] = rootV // Merge the two components.
		mst = append(mst, edge)
		if len(mst) == n-1 {
			break // A spanning tree has exactly n - 1 edges.
		}
	}
	return mst
}

// find returns the root of x's set. Path halving: point each node at its
// grandparent on the way up.
func find(parent []int, x int) int {
	for parent[x] != x {
		parent[x] = parent[parent[x]]
		x = parent[x]
	}
	return x
}

func main() {
	edges := [][3]int{{2, 3, 5}, {0, 1, 1}, {3, 4, 6}, {0, 3, 3}, {1, 3, 2}, {0, 2, 4}}
	mst := kruskal(5, edges)
	fmt.Println(mst) // [[0 1 1] [1 3 2] [0 2 4] [3 4 6]]
	total := 0
	for _, edge := range mst {
		total += edge[2]
	}
	fmt.Println(total) // 13
	// Two separate pieces: only 2 edges for 4 nodes, so there's no spanning tree.
	fmt.Println(kruskal(4, [][3]int{{0, 1, 3}, {2, 3, 2}})) // [[2 3 2] [0 1 3]]
}

using System;
using System.Collections.Generic;
using System.Linq;

public class Solution
{
    // Minimum spanning tree of the nodes 0..n-1, as a list of [u, v, weight]
    // edges. A disconnected graph gets a spanning forest with fewer than n - 1 edges.
    public static List<int[]> Kruskal(int n, int[][] edges)
    {
        int[] parent = Enumerable.Range(0, n).ToArray(); // A small disjoint set: every node starts alone.
        var mst = new List<int[]>();
        // Cheapest edges first. OrderBy is stable, so equal weights keep their input order.
        foreach (int[] edge in edges.OrderBy(e => e[2]))
        {
            int rootU = Find(parent, edge[0]), rootV = Find(parent, edge[1]);
            if (rootU == rootV) continue; // Already connected: this edge would close a cycle.
            parent[rootU] = rootV; // Merge the two components.
            mst.Add(edge);
            if (mst.Count == n - 1) break; // A spanning tree has exactly n - 1 edges.
        }
        return mst;
    }

    // Root of x's set. Path halving: point each node at its grandparent on the way up.
    static int Find(int[] parent, int x)
    {
        while (parent[x] != x)
        {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    public static void Main()
    {
        int[][] edges = { new[] { 2, 3, 5 }, new[] { 0, 1, 1 }, new[] { 3, 4, 6 }, new[] { 0, 3, 3 }, new[] { 1, 3, 2 }, new[] { 0, 2, 4 } };
        var mst = Kruskal(5, edges);
        Console.WriteLine(Format(mst)); // [[0, 1, 1], [1, 3, 2], [0, 2, 4], [3, 4, 6]]
        Console.WriteLine(mst.Sum(edge => edge[2])); // 13
        // Two separate pieces: only 2 edges for 4 nodes, so there's no spanning tree.
        Console.WriteLine(Format(Kruskal(4, new[] { new[] { 0, 1, 3 }, new[] { 2, 3, 2 } }))); // [[2, 3, 2], [0, 1, 3]]
    }

    static string Format(List<int[]> edges) =>
        "[" + string.Join(", ", edges.Select(e => "[" + string.Join(", ", e) + "]")) + "]";
}

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;

public class Solution {
    // Minimum spanning tree of the nodes 0..n-1, as a list of [u, v, weight]
    // edges. A disconnected graph gets a spanning forest with fewer than n - 1 edges.
    static List<int[]> kruskal(int n, int[][] edges) {
        int[] parent = new int[n]; // A small disjoint set: every node starts alone.
        for (int i = 0; i < n; i++) parent[i] = i;
        List<int[]> mst = new ArrayList<>();
        // Cheapest edges first. The sort is stable, so equal weights keep their input order.
        int[][] sorted = edges.clone();
        Arrays.sort(sorted, Comparator.comparingInt(e -> e[2]));
        for (int[] edge : sorted) {
            int rootU = find(parent, edge[0]), rootV = find(parent, edge[1]);
            if (rootU == rootV) continue; // Already connected: this edge would close a cycle.
            parent[rootU] = rootV; // Merge the two components.
            mst.add(edge);
            if (mst.size() == n - 1) break; // A spanning tree has exactly n - 1 edges.
        }
        return mst;
    }

    // Root of x's set. Path halving: point each node at its grandparent on the way up.
    static int find(int[] parent, int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    public static void main(String[] args) {
        int[][] edges = {{2, 3, 5}, {0, 1, 1}, {3, 4, 6}, {0, 3, 3}, {1, 3, 2}, {0, 2, 4}};
        List<int[]> mst = kruskal(5, edges);
        System.out.println(format(mst)); // [[0, 1, 1], [1, 3, 2], [0, 2, 4], [3, 4, 6]]
        System.out.println(mst.stream().mapToInt(edge -> edge[2]).sum()); // 13
        // Two separate pieces: only 2 edges for 4 nodes, so there's no spanning tree.
        System.out.println(format(kruskal(4, new int[][] {{0, 1, 3}, {2, 3, 2}}))); // [[2, 3, 2], [0, 1, 3]]
    }

    static String format(List<int[]> edges) {
        return edges.stream().map(Arrays::toString).collect(Collectors.joining(", ", "[", "]"));
    }
}