Find the Number of Distinct Colors Among the Balls
Source: leetcode
Difficulty: Medium
Tags: array, hash-table, simulation

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 colors among the balls.


Return an array `result` of length `n`, where `result[i]` denotes the number of 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:**


![](https://assets.leetcode.com/uploads/2024/04/17/ezgifcom-crop.gif)


- After query 0, ball 1 has color 4\.
- After query 1, ball 1 has color 4, and ball 2 has color 5\.
- After query 2, ball 1 has color 3, and ball 2 has color 5\.
- After query 3, ball 1 has color 3, ball 2 has color 5, and ball 3 has color 4\.



**Example 2:**



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


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


**Explanation:**


**![](https://assets.leetcode.com/uploads/2024/04/17/ezgifcom-crop2.gif)**


- After query 0, ball 0 has color 1\.
- After query 1, ball 0 has color 1, and ball 1 has color 2\.
- After query 2, ball 0 has color 1, and balls 1 and 2 have color 2\.
- After query 3, ball 0 has color 1, balls 1 and 2 have color 2, and ball 3 has color 4\.
- After query 4, ball 0 has color 1, balls 1 and 2 have color 2, ball 3 has color 4, and ball 4 has color 5\.



 


**Constraints:**


- `1 <= limit <= 109`
- `1 <= n == queries.length <= 105`
- `queries[i].length == 2`
- `0 <= queries[i][0] <= limit`
- `1 <= queries[i][1] <= 109`




C++ solution:
class Solution {
public:
    Solution(){
        ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    }
    vector<int> queryResults(int limit, vector<vector<int>>& queries) {
        unordered_map<int,int> ball_color;
        unordered_map<int,int> cnt_color;
        vector<int> res;
        for(auto& q : queries){
            if(ball_color.count(q[0])){
                int cur_ball_color = ball_color[q[0]];
                cnt_color[cur_ball_color]--;
                if(cnt_color[cur_ball_color] == 0) cnt_color.erase(cur_ball_color);
                
            }
            cnt_color[q[1]]++;
            ball_color[q[0]] = q[1];
            res.push_back(cnt_color.size());
        }
        return res;
    }
};

Go solution:
func queryResults(limit int, queries [][]int) []int {
    ballColors := make(map[int]int)
    colorCount := make(map[int]int)
    result := make([]int, len(queries))
    distinctColors := 0

    for i, query := range queries {
        ball := query[0]
        color := query[1]

        if prevColor, exists := ballColors[ball]; exists {
            colorCount[prevColor]--
            if colorCount[prevColor] == 0 {
                distinctColors--
            }
        }

        ballColors[ball] = color
        colorCount[color]++
        if colorCount[color] == 1 {
            distinctColors++
        }

        result[i] = distinctColors
    }

    return result
}

Java solution:
class Solution {
    public int[] queryResults(int limit, int[][] queries) {
        // Map to track the current color of each ball
        HashMap<Integer, Integer> ballColorMap = new HashMap<>();
        // Map to track the count of each color
        HashMap<Integer, Integer> colorCountMap = new HashMap<>();
        // Result array to store the number of distinct colors after each query
        int[] result = new int[queries.length];
        
        for (int i = 0; i < queries.length; i++) {
            int ball = queries[i][0];
            int color = queries[i][1];
            
            if (ballColorMap.containsKey(ball)) { // Ball is already colored
                int existingColor = ballColorMap.get(ball); // Get existing color
                // Update the count of the existing color
                colorCountMap.put(existingColor, colorCountMap.get(existingColor) - 1);
                if (colorCountMap.get(existingColor) == 0) {
                    colorCountMap.remove(existingColor);
                }
                // Update the ball's color
                ballColorMap.put(ball, color);
                // Update the count of the new color
                colorCountMap.put(color, colorCountMap.getOrDefault(color, 0) + 1);
            } else { // Ball is not colored
                // Set the ball's color
                ballColorMap.put(ball, color);
                // Update the count of the new color
                colorCountMap.put(color, colorCountMap.getOrDefault(color, 0) + 1);
            }
            // Store the number of distinct colors
            result[i] = colorCountMap.size();
        }
        
        return result;
    }
}

JavaScript solution:
/**
 * @param {number} limit
 * @param {number[][]} queries
 * @return {number[]}
 */
var queryResults = function(limit, queries) {
    let ballColors = new Map();
    let colorCount = new Map();
    let result = [];
    let distinctColors = 0;

    for (let [ball, color] of queries) {
        if (ballColors.has(ball)) {
            let prevColor = ballColors.get(ball);
            colorCount.set(prevColor, colorCount.get(prevColor) - 1);
            if (colorCount.get(prevColor) === 0) {
                distinctColors--;
            }
        }

        ballColors.set(ball, color);
        colorCount.set(color, (colorCount.get(color) || 0) + 1);
        if (colorCount.get(color) === 1) {
            distinctColors++;
        }

        result.push(distinctColors);
    }

    return result;
};

Python solution:
class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        res = []
        ball2color = {}
        color2cnt = {}
        for ball, color in queries:
            if ball in ball2color:
                old_color = ball2color[ball]
                color2cnt[old_color] -= 1
                if color2cnt[old_color] == 0:
                    del color2cnt[old_color]
            ball2color[ball] = color
            if color not in color2cnt:
                color2cnt[color] = 0
            color2cnt[color] += 1
            res.append(len(color2cnt))
        return res

Ruby solution:
def query_results(limit, queries)
  ball_color = {}
  color_count = Hash.new(0)
  result = []
  queries.each do |ball, new_color|
    if ball_color.key?(ball)
      old_color = ball_color[ball]
      color_count[old_color] -= 1
      color_count.delete(old_color) if color_count[old_color] == 0
    end
    ball_color[ball] = new_color
    color_count[new_color] += 1
    result << color_count.size
  end
  result
end