LeetCode/go/maximum-number-of-k-divisible-components.go

package main

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

func maxKDivisibleComponents(n int, edges [][]int, values []int, k int) int {
	if n < 2 {
		return 1
	}

	// create empty data structures
	graph := make([][]int, n)
	for i, _ := range graph {
		graph[i] = make([]int, 0)
	}
	inDegrees := make([]int, n)

	// initialize the graph and degrees
	for _, edge := range edges {
		u, v := edge[0], edge[1]

		graph[u] = append(graph[u], v)
		inDegrees[v]++

		graph[v] = append(graph[v], u)
		inDegrees[u]++
	}

	// Start with leafs
	q := aq.New[int]()
	for u, deg := range inDegrees {
		if deg == 1 {
			q.Enqueue(u)
		}
	}

	lValues := make([]int64, n)
	for i, value := range values {
		lValues[i] = int64(value)
	}

	components := 0
	for u, ok := q.Dequeue(); ok; u, ok = q.Dequeue() {
		inDegrees[u]--

		value := int64(0)
		if lValues[u]%int64(k) == 0 {
			components++
		} else {
			value = int64(lValues[u])
		}

		for _, v := range graph[u] {
			if inDegrees[v] == 0 {
				continue
			}

			inDegrees[v]--
			if inDegrees[v] == 1 {
				// if a new leaf
				q.Enqueue(v)
			}

			lValues[v] += value
		}
	}

	return components
}