leetcode-go/hard/685.redundant-connection-ii.go

func findRedundantDirectedConnection(edges [][]int) []int {
	n := len(edges)
	parent := make([]int, n+1)
	rank := make([]int, n+1)
	var e1, e2 []int
	for _, edge := range edges {
		beg, end := edge[0], edge[1]
		if parent[end] != 0 { // Has two parents, one of it is invalid.
			e1 = []int{parent[end], end}
			e2 = []int{beg, end}
			edge[0], edge[1] = 0, 0 // Delete the later edge.
		}
		parent[end] = beg
	}
	parent = make([]int, n+1) // Clear parent for UF.
	for _, edge := range edges {
		beg, end := edge[0], edge[1]
		if beg == 0 || end == 0 {
			continue
		}
		if parent[beg] == 0 {
			parent[beg] = beg
		}
		if parent[end] == 0 {
			parent[end] = end
		}
		i1, i2 := find(parent, beg), find(parent, end)
		if i1 == i2 { // Find loop, e1 or edge is the last of the loop.
			if len(e1) == 0 {
				return edge
			}
			return e1
		}
		if rank[i1] < rank[i2] {
			parent[i1] = i2
		} else {
			parent[i2] = i1
			if rank[i1] == rank[i2] {
				rank[i1]++
			}
		}
	}
	return e2 // Loop not found (after delete e2), so return e2.
}

func find(parent []int, id int) int {
	if parent[id] == id {
		return id
	}
	r := find(parent, parent[id])
	parent[id] = r
	return r
}