diff --git a/go/maximum-number-of-k-divisible-components.go b/go/maximum-number-of-k-divisible-components.go new file mode 100644 index 0000000..a2babfa --- /dev/null +++ b/go/maximum-number-of-k-divisible-components.go @@ -0,0 +1,70 @@ +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 +}