package algo

type Subset[T comparable] struct {
	Parent T
	Rank   int
}

// find is a recursive function that finds the root of the subset that element 'e' belongs to.
// It uses path compression technique to optimize future find operations.
// The subsets parameter is a map that stores the subsets with their parent information.
// The e parameter is the element for which the root is to be found.
// The function returns the root of the subset that 'e' belongs to.
// see https://ondrej-kvasnovsky-2.gitbook.io/algorithms/finding/union-find-algorithms
// or https://jojozhuang.github.io/algorithm/algorithm-union-find/
func Find[T comparable](subsets map[T]*Subset[T], e T) T {
	// find root and make root as parent of e
	// (path compression)
	if subsets[e].Parent != e {
		subsets[e].Parent = Find(subsets, subsets[e].Parent)
	}
	return subsets[e].Parent
}

// Union merges two subsets in the union-find data structure.
// It takes a map of subsets, and the names of the two subsets to be merged.
// The function finds the root of each subset and performs the union operation based on the rank of the subsets.
func Union[T comparable](subsets map[T]*Subset[T], x, y T) {
	xroot := Find(subsets, x)
	yroot := Find(subsets, y)

	// Attach smaller rank tree under root of high
	// rank tree (Union by Rank)
	if subsets[xroot].Rank < subsets[yroot].Rank {
		subsets[xroot].Parent = yroot
	} else {
		if subsets[xroot].Rank > subsets[yroot].Rank {
			subsets[yroot].Parent = xroot
		} else {
			// If ranks are same, then make one as root and
			// increment its rank by one
			subsets[yroot].Parent = xroot
			subsets[xroot].Rank++
		}
	}
}