LeetCode in Kotlin

2493. Divide Nodes Into the Maximum Number of Groups

Hard

You are given a positive integer n representing the number of nodes in an undirected graph. The nodes are labeled from 1 to n.

You are also given a 2D integer array edges, where edges[i] = [ai, bi] indicates that there is a bidirectional edge between nodes ai and bi. Notice that the given graph may be disconnected.

Divide the nodes of the graph into m groups (1-indexed) such that:

Return the maximum number of groups (i.e., maximum m) into which you can divide the nodes. Return -1 if it is impossible to group the nodes with the given conditions.

Example 1:

Input: n = 6, edges = [[1,2],[1,4],[1,5],[2,6],[2,3],[4,6]]

Output: 4

Explanation: As shown in the image we:

We can see that every edge is satisfied. It can be shown that that if we create a fifth group and move any node from the third or fourth group to it, at least on of the edges will not be satisfied.

Example 2:

Input: n = 3, edges = [[1,2],[2,3],[3,1]]

Output: -1

Explanation: If we add node 1 to the first group, node 2 to the second group, and node 3 to the third group to satisfy the first two edges, we can see that the third edge will not be satisfied. It can be shown that no grouping is possible.

Constraints:

Solution

import java.util.LinkedList
import java.util.Queue

class Solution {
    fun magnificentSets(n: Int, edges: Array<IntArray>): Int {
        val adj: MutableList<MutableList<Int>> = ArrayList()
        val visited = IntArray(n + 1)
        visited.fill(-1)
        for (i in 0..n) {
            adj.add(ArrayList())
        }
        for (edge in edges) {
            adj[edge[0]].add(edge[1])
            adj[edge[1]].add(edge[0])
        }
        val comp = IntArray(n + 1)
        var count = -1
        var ans = 0
        for (i in 1..n) {
            if (visited[i] == -1) {
                count++
                comp[count] = bfs(i, adj, visited, count, n)
                if (comp[count] == -1) {
                    return -1
                }
            } else {
                comp[visited[i]] = Math.max(comp[visited[i]], bfs(i, adj, visited, visited[i], n))
            }
        }
        for (group in comp) {
            ans += group
        }
        return ans
    }

    private fun bfs(start: Int, adj: List<MutableList<Int>>, visited: IntArray, count: Int, n: Int): Int {
        val q: Queue<Int> = LinkedList()
        visited[start] = count
        var ans = 1
        val group = IntArray(n + 1)
        q.add(start)
        group[start] = 1
        while (q.isNotEmpty()) {
            val node = q.remove()
            for (adjN in adj[node]) {
                if (group[adjN] == 0) {
                    visited[adjN] = count
                    group[adjN] = group[node] + 1
                    q.add(adjN)
                    ans = Math.max(ans, group[adjN])
                } else if (Math.abs(group[adjN] - group[node]) != 1) {
                    return -1
                }
            }
        }
        return ans
    }
}