LeetCode/go/divide-nodes-into-the-maximum-number-of-groups.go

package main

import aq "github.com/emirpasic/gods/v2/queues/arrayqueue"

func magnificentSets(n int, edges [][]int) int {
	g := make([][]int, n)
	for u, _ := range g {
		g[u] = make([]int, 0)
	}

	depth := make([]int, n)
	parent := make([]int, n)
	for u, _ := range parent {
		parent[u] = -1
	}

	find := func(t int) int {
		for parent[t] != -1 {
			t = parent[t]
		}
		return t
	}

	union := func(u, v int) {
		su, sv := find(u), find(v)
		if su == sv {
			// same component
			return
		}

		if depth[su] < depth[sv] {
			su, sv = sv, su
		}
		parent[sv] = su

		if depth[su] == depth[sv] {
			depth[su]++
		}
	}

	numberOfGroups := func(src int) int {
		q := aq.New[int]()
		seen := make([]int, n)
		for i, _ := range seen {
			seen[i] = -1
		}

		q.Enqueue(src)
		seen[src] = 0

		deepest := 0
		for !q.Empty() {
			size := q.Size()
			for i := 0; i < size; i++ {
				u, _ := q.Dequeue()
				for _, v := range g[u] {
					if seen[v] == -1 {
						seen[v] = deepest + 1
						q.Enqueue(v)
					} else if seen[v] == deepest {
						return -1
					}
				}
			}
			deepest++
		}

		return deepest
	}

	for _, edge := range edges {
		u, v := edge[0]-1, edge[1]-1
		g[u] = append(g[u], v)
		g[v] = append(g[v], u)
		union(u, v)
	}

	groupsPerComponent := make(map[int]int)
	for u, _ := range g {
		groups := numberOfGroups(u)
		if groups == -1 {
			// invalid
			return -1
		}

		root := find(u)

		current, ok := groupsPerComponent[root]
		if !ok {
			current = 0
		}
		groupsPerComponent[root] = max(current, groups)
	}

	totalGroups := 0
	for _, groups := range groupsPerComponent {
		totalGroups += groups
	}
	return totalGroups
}