LeetCode-in-Java

3160. Find the Number of Distinct Colors Among the Balls

Medium

You are given an integer limit and a 2D array queries of size n x 2.

There are limit + 1 balls with distinct labels in the range [0, limit]. Initially, all balls are uncolored. For every query in queries that is of the form [x, y], you mark ball x with the color y. After each query, you need to find the number of distinct colors among the balls.

Return an array result of length n, where result[i] denotes the number of distinct colors after ith query.

Note that when answering a query, lack of a color will not be considered as a color.

Example 1:

Input: limit = 4, queries = [[1,4],[2,5],[1,3],[3,4]]

Output: [1,2,2,3]

Explanation:

Example 2:

Input: limit = 4, queries = [[0,1],[1,2],[2,2],[3,4],[4,5]]

Output: [1,2,2,3,4]

Explanation:

Constraints:

Solution

import java.util.HashMap;
import java.util.Map;

@SuppressWarnings("java:S1172")
public class Solution {
    public int[] queryResults(int ignoredLimit, int[][] queries) {
        Map<Integer, Integer> ballToColor = new HashMap<>();
        Map<Integer, Integer> colorToCnt = new HashMap<>();
        int[] ret = new int[queries.length];
        for (int i = 0; i < queries.length; i += 1) {
            final int ball = queries[i][0];
            final int color = queries[i][1];
            if (ballToColor.containsKey(ball)) {
                int oldColor = ballToColor.get(ball);
                int oldColorCnt = colorToCnt.get(oldColor);
                if (oldColorCnt >= 2) {
                    colorToCnt.put(oldColor, oldColorCnt - 1);
                } else {
                    colorToCnt.remove(oldColor);
                }
            }
            ballToColor.put(ball, color);
            colorToCnt.put(color, colorToCnt.getOrDefault(color, 0) + 1);
            ret[i] = colorToCnt.size();
        }
        return ret;
    }
}