Find the Number of Distinct Colors Among the Balls

Med

#2807Time: O(n^2), where n is the number of queries. For each of the n queries, we update a map (O(1)) and then iterate through all its current values to count distinct colors. The map can grow up to size `i` at step `i`. The total time is the sum of `i` from 1 to `n`, which is `1 + 2 + ... + n-1 = O(n^2)`.Space: O(n), where n is the number of queries. The `ballToColor` map can store up to `n` entries if all queries are for distinct balls. The `HashSet` used for counting also takes up to O(n) space in the worst case.

Data structures

Array Hash Table

Prompt

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 i^th 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:

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:

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 <= 10⁹
1 <= n == queries.length <= 10⁵
queries[i].length == 2
0 <= queries[i][0] <= limit
1 <= queries[i][1] <= 10⁹

Approaches

2 approaches with complexity analysis and trade-offs.

This approach processes each query one by one. After each query, it determines the current color of every ball that has been mentioned so far and then counts the number of unique colors among them. This is straightforward but inefficient as it re-evaluates the state of the system repeatedly.

Algorithm

Initialize an empty HashMap called ballToColor to store the mapping from a ball's label to its color.
Initialize an integer array result of the same size as queries.
Iterate through the queries array from i = 0 to n-1:
- For the current query [ball, color], update the map: ballToColor.put(ball, color).
- Create a new HashSet from the values of the ballToColor map.
- The size of this HashSet is the number of distinct colors.
- Store this size in result[i].
After the loop, return the result array.

Walkthrough

We maintain a map, ballToColor, to store the current color of each ball. For each query, we update this map. After the update, we find the number of distinct colors by iterating through all the values (colors) in the ballToColor map and adding them to a HashSet. The size of the HashSet gives us the number of distinct colors. This process is repeated for every single query. The main drawback is the recalculation of distinct colors from scratch in every step, which involves iterating over all currently colored balls.

1class Solution {2    public int[] getNumberOfDistinctColors(int limit, int[][] queries) {3        int n = queries.length;4        int[] result = new int[n];5        Map<Integer, Integer> ballToColor = new HashMap<>();6 7        for (int i = 0; i < n; i++) {8            int ball = queries[i][0];9            int color = queries[i][1];10 11            // Update the color of the ball12            ballToColor.put(ball, color);13 14            // Recalculate the number of distinct colors15            Set<Integer> distinctColors = new HashSet<>(ballToColor.values());16            result[i] = distinctColors.size();17        }18 19        return result;20    }21}

Complexity

Time

O(n^2), where n is the number of queries. For each of the n queries, we update a map (O(1)) and then iterate through all its current values to count distinct colors. The map can grow up to size `i` at step `i`. The total time is the sum of `i` from 1 to `n`, which is `1 + 2 + ... + n-1 = O(n^2)`.

Space

O(n), where n is the number of queries. The `ballToColor` map can store up to `n` entries if all queries are for distinct balls. The `HashSet` used for counting also takes up to O(n) space in the worst case.

Trade-offs

Pros

Simple to understand and implement.
Correctly solves the problem.

Cons

Inefficient for large inputs, likely resulting in a 'Time Limit Exceeded' error.
Performs a lot of redundant work by recounting all colors in every step.

This optimized approach avoids recalculation by maintaining the counts of each color in real-time. It uses two maps: one to track the color of each ball (ballToColor) and another to track how many balls have a certain color (colorToCount). This allows for O(1) updates per query, making it highly efficient.

Algorithm

Initialize two empty HashMaps: ballToColor (ball -> color) and colorToCount (color -> count).
Initialize an integer array result of size n.
Iterate through the queries array from i = 0 to n-1:
- Let the current query be [ball, newColor].
- Check if ball exists as a key in ballToColor.
- If it exists:
  - Get its oldColor.
  - If oldColor is the same as newColor, no change in distinct colors, so continue to the next query after recording the current count.
  - Decrement the count of oldColor in colorToCount.
  - If the count of oldColor becomes 0, remove it from colorToCount.
- Update the ball's color in the ballToColor map: ballToColor.put(ball, newColor).
- Increment the count of newColor in colorToCount (using getOrDefault for convenience).
- The number of distinct colors is colorToCount.size(). Store this in result[i].
After the loop, return the result array.

Walkthrough

We use two HashMaps:

ballToColor: Maps a ball's label to its current color.
colorToCount: Maps a color to the number of balls currently painted with that color.

When processing a query [x, y]:

Check if ball x already has a color by looking it up in ballToColor.
If it does (oldColor), we decrement the count of oldColor in colorToCount. If this count drops to zero, it means no ball has this color anymore, so we remove oldColor from colorToCount.
We then update ball x's color to y in ballToColor.
We increment the count of the new color y in colorToCount. If y was not in colorToCount before, it's added with a count of 1.

The number of distinct colors after each query is simply the size of the colorToCount map. This is because this map only holds entries for colors that are currently in use (i.e., have a count > 0).

1class Solution {2    public int[] getNumberOfDistinctColors(int limit, int[][] queries) {3        int n = queries.length;4        int[] result = new int[n];5        Map<Integer, Integer> ballToColor = new HashMap<>();6        Map<Integer, Integer> colorToCount = new HashMap<>();7 8        for (int i = 0; i < n; i++) {9            int ball = queries[i][0];10            int newColor = queries[i][1];11 12            // Check if the ball was already colored13            if (ballToColor.containsKey(ball)) {14                int oldColor = ballToColor.get(ball);15                16                // If the new color is the same as the old one, nothing changes17                if (oldColor == newColor) {18                    result[i] = colorToCount.size();19                    continue;20                }21 22                // Decrement the count of the old color23                int oldCount = colorToCount.get(oldColor);24                if (oldCount == 1) {25                    colorToCount.remove(oldColor);26                } else {27                    colorToCount.put(oldColor, oldCount - 1);28                }29            }30 31            // Update the ball's color32            ballToColor.put(ball, newColor);33 34            // Increment the count of the new color35            colorToCount.put(newColor, colorToCount.getOrDefault(newColor, 0) + 1);36 37            // The number of distinct colors is the size of the colorToCount map38            result[i] = colorToCount.size();39        }40 41        return result;42    }43}

Complexity

Time

O(n), where n is the number of queries. Each query involves a constant number of `HashMap` operations (get, put, remove), which take O(1) time on average. Therefore, the total time complexity is linear with respect to the number of queries.

Space

O(n), where n is the number of queries. In the worst-case scenario, each query could involve a new ball and a new color. This would mean both `ballToColor` and `colorToCount` maps could grow to a size of up to `n`.

Trade-offs

Pros

Highly efficient, with linear time complexity.
Scales well with large numbers of queries.
Optimal solution for the given constraints.

Cons

Slightly more complex to implement than the brute-force approach due to the need to manage two maps and their state correctly.

Solutions

Solution

1class Solution {2public3  int[] queryResults(int limit, int[][] queries) {4    Map<Integer, Integer> g = new HashMap<>();5    Map<Integer, Integer> cnt = new HashMap<>();6    int m = queries.length;7    int[] ans = new int[m];8    for (int i = 0; i < m; ++i) {9      int x = queries[i][0], y = queries[i][1];10      cnt.merge(y, 1, Integer : : sum);11      if (g.containsKey(x) && cnt.merge(g.get(x), -1, Integer : : sum) == 0) {12        cnt.remove(g.get(x));13      }14      g.put(x, y);15      ans[i] = cnt.size();16    }17    return ans;18  }19}

1class Solution {2public:3  vector<int> queryResults(int limit, vector<vector<int>> &queries) {4    unordered_map<int, int> g;5    unordered_map<int, int> cnt;6    vector<int> ans;7    for (auto &q : queries) {8      int x = q[0], y = q[1];9      cnt[y]++;10      if (g.contains(x) && --cnt[g[x]] == 0) {11        cnt.erase(g[x]);12      }13      g[x] = y;14      ans.push_back(cnt.size());15    }16    return ans;17  }18};

1class Solution:2    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]: g = {} cnt = Counter() ans = [] for x, y in queries: cnt[y] += 1 if x in g: cnt[g[x]] -= 1 if cnt[g[x]] == 0: cnt . pop(g[x]) g[x] = y ans . append(len(cnt)) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem