LeetCode/kt/find-the-number-of-distinct-colors-among-the-balls.kt

class Solution {
    fun queryResults(
        limit: Int,
        queries: Array<IntArray>,
    ): IntArray {
        val n = queries.size
        val freqs = mutableMapOf<Int, Int>()
        val balls = mutableMapOf<Int, Int>()

        return queries.map { q ->
            val (ball, color) = q[0] to q[1]

            if (balls.containsKey(ball)) {
                val currentColor = balls[ball]!!
                freqs.put(currentColor, freqs.get(currentColor)!! - 1)

                if (freqs[currentColor] == 0) {
                    freqs.remove(currentColor)
                }
            }

            balls[ball] = color
            freqs[color] = 1 + (freqs.get(color) ?: 0)

            freqs.size
        }.toIntArray()
    }
}