Union-Find

def find_circle_num(is_connected: list[list[int]]) -> int:
    """How many groups of cities are linked, directly or through other cities."""
    n = len(is_connected)
    ds = DisjointSet(n)  # Every city starts as its own province.
    for i in range(n):
        # The matrix is symmetric, so the upper triangle has every road once.
        for j in range(i + 1, n):
            if is_connected[i][j] == 1:
                ds.union(i, j)
    return ds.count


class DisjointSet:
    """The disjoint set from the template: path compression plus union by size."""

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n
        self.count = n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a: int, b: int) -> bool:
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return False
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]
        self.count -= 1  # Two provinces just became one.
        return True


if __name__ == "__main__":
    print(find_circle_num([[1, 1, 0], [1, 1, 0], [0, 0, 1]]))  # 2
    print(find_circle_num([[1, 0, 0], [0, 1, 0], [0, 0, 1]]))  # 3
    print(find_circle_num([[1, 0, 0, 1], [0, 1, 1, 0], [0, 1, 1, 1], [1, 0, 1, 1]]))  # 1

// How many groups of cities are linked, directly or through other cities.
export function findCircleNum(isConnected: number[][]): number {
  const n = isConnected.length;
  const ds = new DisjointSet(n); // Every city starts as its own province.
  for (let i = 0; i < n; i++) {
    // The matrix is symmetric, so the upper triangle has every road once.
    for (let j = i + 1; j < n; j++) {
      if (isConnected[i][j] === 1) ds.union(i, j);
    }
  }
  return ds.count;
}

// The disjoint set from the template: path compression plus union by size.
class DisjointSet {
  private parent: number[];
  private size: number[];
  count: number;

  constructor(n: number) {
    this.parent = Array.from({ length: n }, (_, i) => i);
    this.size = new Array<number>(n).fill(1);
    this.count = n;
  }

  find(x: number): number {
    if (this.parent[x] !== x) this.parent[x] = this.find(this.parent[x]);
    return this.parent[x];
  }

  union(a: number, b: number): boolean {
    let rootA = this.find(a);
    let rootB = this.find(b);
    if (rootA === rootB) return false;
    if (this.size[rootA] < this.size[rootB]) [rootA, rootB] = [rootB, rootA];
    this.parent[rootB] = rootA;
    this.size[rootA] += this.size[rootB];
    this.count--; // Two provinces just became one.
    return true;
  }
}

console.log(findCircleNum([[1, 1, 0], [1, 1, 0], [0, 0, 1]])); // 2
console.log(findCircleNum([[1, 0, 0], [0, 1, 0], [0, 0, 1]])); // 3
console.log(findCircleNum([[1, 0, 0, 1], [0, 1, 1, 0], [0, 1, 1, 1], [1, 0, 1, 1]])); // 1

package main

import "fmt"

// findCircleNum returns how many groups of cities are linked, directly or
// through other cities.
func findCircleNum(isConnected [][]int) int {
	n := len(isConnected)
	ds := newDisjointSet(n) // Every city starts as its own province.
	for i := range n {
		// The matrix is symmetric, so the upper triangle has every road once.
		for j := i + 1; j < n; j++ {
			if isConnected[i][j] == 1 {
				ds.union(i, j)
			}
		}
	}
	return ds.count
}

// disjointSet is the disjoint set from the template: path compression plus
// union by size.
type disjointSet struct {
	parent []int
	size   []int
	count  int
}

func newDisjointSet(n int) *disjointSet {
	ds := &disjointSet{parent: make([]int, n), size: make([]int, n), count: n}
	for i := range n {
		ds.parent[i] = i
		ds.size[i] = 1
	}
	return ds
}

func (ds *disjointSet) find(x int) int {
	if ds.parent[x] != x {
		ds.parent[x] = ds.find(ds.parent[x])
	}
	return ds.parent[x]
}

func (ds *disjointSet) union(a, b int) bool {
	rootA, rootB := ds.find(a), ds.find(b)
	if rootA == rootB {
		return false
	}
	if ds.size[rootA] < ds.size[rootB] {
		rootA, rootB = rootB, rootA
	}
	ds.parent[rootB] = rootA
	ds.size[rootA] += ds.size[rootB]
	ds.count-- // Two provinces just became one.
	return true
}

func main() {
	fmt.Println(findCircleNum([][]int{{1, 1, 0}, {1, 1, 0}, {0, 0, 1}}))                        // 2
	fmt.Println(findCircleNum([][]int{{1, 0, 0}, {0, 1, 0}, {0, 0, 1}}))                        // 3
	fmt.Println(findCircleNum([][]int{{1, 0, 0, 1}, {0, 1, 1, 0}, {0, 1, 1, 1}, {1, 0, 1, 1}})) // 1
}

using System;

public class Solution
{
    // How many groups of cities are linked, directly or through other cities.
    public static int FindCircleNum(int[][] isConnected)
    {
        int n = isConnected.Length;
        var ds = new DisjointSet(n); // Every city starts as its own province.
        for (int i = 0; i < n; i++)
        {
            // The matrix is symmetric, so the upper triangle has every road once.
            for (int j = i + 1; j < n; j++)
            {
                if (isConnected[i][j] == 1) ds.Union(i, j);
            }
        }
        return ds.Count;
    }

    // The disjoint set from the template: path compression plus union by size.
    public class DisjointSet
    {
        readonly int[] parent;
        readonly int[] size;
        public int Count { get; private set; }

        public DisjointSet(int n)
        {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++)
            {
                parent[i] = i;
                size[i] = 1;
            }
            Count = n;
        }

        public int Find(int x)
        {
            if (parent[x] != x) parent[x] = Find(parent[x]);
            return parent[x];
        }

        public bool Union(int a, int b)
        {
            int rootA = Find(a), rootB = Find(b);
            if (rootA == rootB) return false;
            if (size[rootA] < size[rootB]) (rootA, rootB) = (rootB, rootA);
            parent[rootB] = rootA;
            size[rootA] += size[rootB];
            Count--; // Two provinces just became one.
            return true;
        }
    }

    public static void Main()
    {
        Console.WriteLine(FindCircleNum(new[] { new[] { 1, 1, 0 }, new[] { 1, 1, 0 }, new[] { 0, 0, 1 } })); // 2
        Console.WriteLine(FindCircleNum(new[] { new[] { 1, 0, 0 }, new[] { 0, 1, 0 }, new[] { 0, 0, 1 } })); // 3
        Console.WriteLine(FindCircleNum(new[] { new[] { 1, 0, 0, 1 }, new[] { 0, 1, 1, 0 }, new[] { 0, 1, 1, 1 }, new[] { 1, 0, 1, 1 } })); // 1
    }
}

public class Solution {
    // How many groups of cities are linked, directly or through other cities.
    static int findCircleNum(int[][] isConnected) {
        int n = isConnected.length;
        DisjointSet ds = new DisjointSet(n); // Every city starts as its own province.
        for (int i = 0; i < n; i++) {
            // The matrix is symmetric, so the upper triangle has every road once.
            for (int j = i + 1; j < n; j++) {
                if (isConnected[i][j] == 1) ds.union(i, j);
            }
        }
        return ds.count;
    }

    // The disjoint set from the template: path compression plus union by size.
    static class DisjointSet {
        private final int[] parent;
        private final int[] size;
        int count;

        DisjointSet(int n) {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                size[i] = 1;
            }
            count = n;
        }

        int find(int x) {
            if (parent[x] != x) parent[x] = find(parent[x]);
            return parent[x];
        }

        boolean union(int a, int b) {
            int rootA = find(a), rootB = find(b);
            if (rootA == rootB) return false;
            if (size[rootA] < size[rootB]) {
                int swap = rootA;
                rootA = rootB;
                rootB = swap;
            }
            parent[rootB] = rootA;
            size[rootA] += size[rootB];
            count--; // Two provinces just became one.
            return true;
        }
    }

    public static void main(String[] args) {
        System.out.println(findCircleNum(new int[][] {{1, 1, 0}, {1, 1, 0}, {0, 0, 1}})); // 2
        System.out.println(findCircleNum(new int[][] {{1, 0, 0}, {0, 1, 0}, {0, 0, 1}})); // 3
        System.out.println(findCircleNum(new int[][] {{1, 0, 0, 1}, {0, 1, 1, 0}, {0, 1, 1, 1}, {1, 0, 1, 1}})); // 1
    }
}

def find_redundant_connection(edges: list[list[int]]) -> list[int]:
    """The edge to remove from a tree with one extra edge: the one that closes
    the cycle, which is the cycle edge that appears last in `edges`."""
    ds = DisjointSet(len(edges) + 1)  # Nodes run 1..n, and n is len(edges).
    for a, b in edges:
        # Both ends already connected: this edge closes the cycle.
        if not ds.union(a, b):
            return [a, b]
    return []


class DisjointSet:
    """The disjoint set from the template: path compression plus union by size."""

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a: int, b: int) -> bool:
        """Returns False when a and b were already in the same set."""
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return False
        if self.size[root_a] < self.size[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        self.size[root_a] += self.size[root_b]
        return True


if __name__ == "__main__":
    print(find_redundant_connection([[1, 2], [1, 3], [2, 3]]))  # [2, 3]
    print(find_redundant_connection([[1, 2], [2, 3], [3, 4], [1, 4], [1, 5]]))  # [1, 4]
    print(find_redundant_connection([[1, 2], [2, 3], [1, 3], [3, 4]]))  # [1, 3]

// The edge to remove from a tree with one extra edge: the one that closes
// the cycle, which is the cycle edge that appears last in `edges`.
export function findRedundantConnection(edges: number[][]): number[] {
  const ds = new DisjointSet(edges.length + 1); // Nodes run 1..n, and n is edges.length.
  for (const [a, b] of edges) {
    // Both ends already connected: this edge closes the cycle.
    if (!ds.union(a, b)) return [a, b];
  }
  return [];
}

// The disjoint set from the template: path compression plus union by size.
class DisjointSet {
  private parent: number[];
  private size: number[];

  constructor(n: number) {
    this.parent = Array.from({ length: n }, (_, i) => i);
    this.size = new Array<number>(n).fill(1);
  }

  find(x: number): number {
    if (this.parent[x] !== x) this.parent[x] = this.find(this.parent[x]);
    return this.parent[x];
  }

  // Returns false when a and b were already in the same set.
  union(a: number, b: number): boolean {
    let rootA = this.find(a);
    let rootB = this.find(b);
    if (rootA === rootB) return false;
    if (this.size[rootA] < this.size[rootB]) [rootA, rootB] = [rootB, rootA];
    this.parent[rootB] = rootA;
    this.size[rootA] += this.size[rootB];
    return true;
  }
}

console.log(JSON.stringify(findRedundantConnection([[1, 2], [1, 3], [2, 3]]))); // [2,3]
console.log(JSON.stringify(findRedundantConnection([[1, 2], [2, 3], [3, 4], [1, 4], [1, 5]]))); // [1,4]
console.log(JSON.stringify(findRedundantConnection([[1, 2], [2, 3], [1, 3], [3, 4]]))); // [1,3]

package main

import "fmt"

// findRedundantConnection returns the edge to remove from a tree with one
// extra edge: the one that closes the cycle, which is the cycle edge that
// appears last in edges.
func findRedundantConnection(edges [][]int) []int {
	ds := newDisjointSet(len(edges) + 1) // Nodes run 1..n, and n is len(edges).
	for _, e := range edges {
		// Both ends already connected: this edge closes the cycle.
		if !ds.union(e[0], e[1]) {
			return e
		}
	}
	return nil
}

// disjointSet is the disjoint set from the template: path compression plus
// union by size.
type disjointSet struct {
	parent []int
	size   []int
}

func newDisjointSet(n int) *disjointSet {
	ds := &disjointSet{parent: make([]int, n), size: make([]int, n)}
	for i := range n {
		ds.parent[i] = i
		ds.size[i] = 1
	}
	return ds
}

func (ds *disjointSet) find(x int) int {
	if ds.parent[x] != x {
		ds.parent[x] = ds.find(ds.parent[x])
	}
	return ds.parent[x]
}

// union returns false when a and b were already in the same set.
func (ds *disjointSet) union(a, b int) bool {
	rootA, rootB := ds.find(a), ds.find(b)
	if rootA == rootB {
		return false
	}
	if ds.size[rootA] < ds.size[rootB] {
		rootA, rootB = rootB, rootA
	}
	ds.parent[rootB] = rootA
	ds.size[rootA] += ds.size[rootB]
	return true
}

func main() {
	fmt.Println(findRedundantConnection([][]int{{1, 2}, {1, 3}, {2, 3}}))                 // [2 3]
	fmt.Println(findRedundantConnection([][]int{{1, 2}, {2, 3}, {3, 4}, {1, 4}, {1, 5}})) // [1 4]
	fmt.Println(findRedundantConnection([][]int{{1, 2}, {2, 3}, {1, 3}, {3, 4}}))         // [1 3]
}

using System;

public class Solution
{
    // The edge to remove from a tree with one extra edge: the one that closes
    // the cycle, which is the cycle edge that appears last in edges.
    public static int[] FindRedundantConnection(int[][] edges)
    {
        var ds = new DisjointSet(edges.Length + 1); // Nodes run 1..n, and n is edges.Length.
        foreach (int[] edge in edges)
        {
            // Both ends already connected: this edge closes the cycle.
            if (!ds.Union(edge[0], edge[1])) return edge;
        }
        return Array.Empty<int>();
    }

    // The disjoint set from the template: path compression plus union by size.
    public class DisjointSet
    {
        readonly int[] parent;
        readonly int[] size;

        public DisjointSet(int n)
        {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++)
            {
                parent[i] = i;
                size[i] = 1;
            }
        }

        public int Find(int x)
        {
            if (parent[x] != x) parent[x] = Find(parent[x]);
            return parent[x];
        }

        // Returns false when a and b were already in the same set.
        public bool Union(int a, int b)
        {
            int rootA = Find(a), rootB = Find(b);
            if (rootA == rootB) return false;
            if (size[rootA] < size[rootB]) (rootA, rootB) = (rootB, rootA);
            parent[rootB] = rootA;
            size[rootA] += size[rootB];
            return true;
        }
    }

    public static void Main()
    {
        Console.WriteLine(Format(FindRedundantConnection(new[] { new[] { 1, 2 }, new[] { 1, 3 }, new[] { 2, 3 } }))); // [2, 3]
        Console.WriteLine(Format(FindRedundantConnection(new[] { new[] { 1, 2 }, new[] { 2, 3 }, new[] { 3, 4 }, new[] { 1, 4 }, new[] { 1, 5 } }))); // [1, 4]
        Console.WriteLine(Format(FindRedundantConnection(new[] { new[] { 1, 2 }, new[] { 2, 3 }, new[] { 1, 3 }, new[] { 3, 4 } }))); // [1, 3]
    }

    static string Format(int[] values) => "[" + string.Join(", ", values) + "]";
}

import java.util.Arrays;

public class Solution {
    // The edge to remove from a tree with one extra edge: the one that closes
    // the cycle, which is the cycle edge that appears last in edges.
    static int[] findRedundantConnection(int[][] edges) {
        DisjointSet ds = new DisjointSet(edges.length + 1); // Nodes run 1..n, and n is edges.length.
        for (int[] edge : edges) {
            // Both ends already connected: this edge closes the cycle.
            if (!ds.union(edge[0], edge[1])) return edge;
        }
        return new int[0];
    }

    // The disjoint set from the template: path compression plus union by size.
    static class DisjointSet {
        private final int[] parent;
        private final int[] size;

        DisjointSet(int n) {
            parent = new int[n];
            size = new int[n];
            for (int i = 0; i < n; i++) {
                parent[i] = i;
                size[i] = 1;
            }
        }

        int find(int x) {
            if (parent[x] != x) parent[x] = find(parent[x]);
            return parent[x];
        }

        // Returns false when a and b were already in the same set.
        boolean union(int a, int b) {
            int rootA = find(a), rootB = find(b);
            if (rootA == rootB) return false;
            if (size[rootA] < size[rootB]) {
                int swap = rootA;
                rootA = rootB;
                rootB = swap;
            }
            parent[rootB] = rootA;
            size[rootA] += size[rootB];
            return true;
        }
    }

    public static void main(String[] args) {
        System.out.println(Arrays.toString(findRedundantConnection(new int[][] {{1, 2}, {1, 3}, {2, 3}}))); // [2, 3]
        System.out.println(Arrays.toString(findRedundantConnection(new int[][] {{1, 2}, {2, 3}, {3, 4}, {1, 4}, {1, 5}}))); // [1, 4]
        System.out.println(Arrays.toString(findRedundantConnection(new int[][] {{1, 2}, {2, 3}, {1, 3}, {3, 4}}))); // [1, 3]
    }
}