leetcode-go/leetcode/medium/684.redundant-connection.go

type unionFind struct {
	id []int
	sz []int
}

func newUF(k int) (uf unionFind) {
	uf.id = make([]int, k+1)
	uf.sz = make([]int, k+1)
	for i := range uf.id {
		uf.id[i] = i
	}
	return
}

func (uf unionFind) union(x, y int) bool {
	x, y = uf.find(x), uf.find(y)
	if x == y {
		return false
	}
	if uf.sz[x] < uf.sz[y] {
		uf.id[x] = y
	} else {
		uf.id[y] = x
		if uf.sz[x] == uf.sz[y] {
			uf.sz[x]++
		}
	}
	return true
}

func (uf unionFind) find(x int) int {
	if uf.id[x] == x {
		return x
	}
	i := uf.find(uf.id[x])
	uf.id[x] = i
	return i
}

func findRedundantConnection(edges [][]int) []int {
	n := len(edges)
	uf := newUF(n)
	for _, edge := range edges {
		if !uf.union(edge[0], edge[1]) {
			return edge
		}
	}
	return []int{}
}