Disjoint Set Union (Union-Find) with Path Compression

package main

import "fmt"

type DSU struct {
	parent []int
	rank   []int
}

func NewDSU(n int) *DSU {
	parent := make([]int, n)
	rank := make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
	}
	return &DSU{parent: parent, rank: rank}
}

func (d *DSU) Find(x int) int {
	if d.parent[x] != x {
		d.parent[x] = d.Find(d.parent[x])
	}
	return d.parent[x]
}

func (d *DSU) Union(a, b int) bool {
	ra, rb := d.Find(a), d.Find(b)
	if ra == rb {
		return false
	}
	if d.rank[ra] < d.rank[rb] {
		ra, rb = rb, ra
	}
	d.parent[rb] = ra
	if d.rank[ra] == d.rank[rb] {
		d.rank[ra]++
	}
	return true
}

func (d *DSU) Connected(a, b int) bool {
	return d.Find(a) == d.Find(b)
}

func main() {
	d := NewDSU(6)
	d.Union(0, 1)
	d.Union(2, 3)
	d.Union(1, 2)

	fmt.Println("0 connected to 3?", d.Connected(0, 3))
	fmt.Println("4 connected to 5?", d.Connected(4, 5))
}