LeetCode 3160 - Find the Number of Distinct Colors Among the Balls

LeetCode Problem 3160

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Simulation

Solution

Problem Understanding

The problem gives us limit + 1 balls labeled from 0 to limit. Initially, none of the balls have a color assigned. We then process a sequence of queries, where each query is of the form [x, y]. This means ball x should now be painted with color y.

The important detail is that repainting is allowed. If a ball already has a color and receives another query, its old color is replaced by the new one. After processing each query, we must determine how many distinct colors currently exist among all colored balls.

Uncolored balls do not count as having a color. Only colors that are actively assigned to at least one ball should be counted.

For example, suppose ball 1 is colored with 4, and later recolored with 3. If no other ball still has color 4, then color 4 disappears entirely from the system and should no longer be counted.

The input consists of:

• limit, which defines the valid ball labels from 0 to limit
• queries, where each query updates the color of one ball

The output is an array where result[i] is the number of distinct colors after processing the i-th query.

The constraints are very important:

• queries.length can be as large as 10^5
• Ball labels can be very large, up to 10^9
• Colors can also be very large, up to 10^9

These constraints immediately tell us that we cannot use large arrays indexed directly by ball label or color value. Even though there are potentially billions of labels, only up to 10^5 queries exist, so at most 10^5 balls will ever actually appear. This strongly suggests using hash maps instead of arrays.

Several edge cases are important:

• A ball may be recolored multiple times.
• Multiple balls may share the same color.
• Recoloring a ball to an already existing color should not increase the distinct color count.
• Removing the last ball using a color should decrease the distinct color count.
• A query may repaint a ball with the same color it already has.

Approaches

Brute Force Approach

A straightforward solution is to maintain the current color of every ball and, after each query, recompute the number of distinct colors from scratch.

We could store the current color assignment in a hash map:

• ball_to_color[ball] = current color

After each query:

1. Update the ball's color.
2. Iterate through all colored balls.
3. Insert their colors into a set.
4. Return the size of the set.

This approach is correct because the set naturally removes duplicates, leaving only distinct colors.

However, it is too slow. If there are n queries, and after every query we scan up to n balls, the total complexity becomes O(n^2) in the worst case.

With n = 10^5, this would require up to 10^10 operations, which is far beyond acceptable limits.

Optimal Approach

The key observation is that we do not need to recompute distinct colors from scratch after every query.

Instead, we can maintain the frequency of each color incrementally.

We use two hash maps:

• ball_to_color tracks the current color of each ball.
• color_count tracks how many balls currently use each color.

When processing a query:

1. If the ball already has a color:

• Decrease the frequency of the old color.
• If the frequency becomes zero, remove that color entirely.

2. Assign the new color to the ball.
3. Increase the frequency of the new color.
4. The number of distinct colors is simply the number of keys in color_count.

This works because color_count always reflects exactly how many balls currently use each color.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(n²)	O(n)	Recomputes distinct colors after every query
Optimal	O(n)	O(n)	Maintains color frequencies incrementally

Algorithm Walkthrough

1. Create a hash map called ball_to_color.

This stores the current color assigned to each ball. If a ball does not exist in the map, it is currently uncolored. 2. Create another hash map called color_count.

This stores how many balls currently use each color. 3. Create an empty result array.

After processing each query, we append the current number of distinct colors. 4. Process each query [ball, new_color].

For every query, first check whether the ball already has a color assigned. 5. If the ball already has an old color:

Retrieve the old color from ball_to_color.

Decrease its count in color_count.

If the count becomes zero, remove the color from the map entirely. This is important because only colors used by at least one ball should count as distinct. 6. Assign the new color to the ball.

Update ball_to_color[ball] = new_color. 7. Increase the frequency of the new color.

Increment color_count[new_color]. 8. Record the answer.

The number of distinct colors is simply len(color_count) because every key in the map represents an active color. 9. Continue until all queries are processed.

Why it works

The algorithm maintains a precise invariant:

• ball_to_color always stores the current color of every colored ball.
• color_count[color] always equals the number of balls currently painted with that color.

Whenever a ball changes color, we remove one occurrence of the old color and add one occurrence of the new color. If a color's frequency becomes zero, it is removed entirely.

Therefore, the number of keys in color_count is always exactly the number of distinct colors currently present.

Python Solution

from typing import List
from collections import defaultdict

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        ball_to_color = {}
        color_count = defaultdict(int)

        result = []

        for ball, new_color in queries:
            # Remove the old color if the ball was already colored
            if ball in ball_to_color:
                old_color = ball_to_color[ball]

                color_count[old_color] -= 1

                if color_count[old_color] == 0:
                    del color_count[old_color]

            # Assign the new color
            ball_to_color[ball] = new_color

            # Add the new color frequency
            color_count[new_color] += 1

            # Number of distinct colors
            result.append(len(color_count))

        return result

The implementation directly follows the algorithm described earlier.

ball_to_color stores the latest color for each ball. Since only queried balls matter, a dictionary is sufficient and memory efficient.

color_count tracks how many balls currently use each color. A defaultdict(int) simplifies increment and decrement operations because missing keys automatically start at zero.

For every query, the code first checks whether the ball already has an assigned color. If so, the old color frequency is decreased. When that frequency reaches zero, the color is removed from the dictionary entirely so it no longer contributes to the distinct color count.

Next, the new color is assigned to the ball, and its frequency is incremented.

Finally, len(color_count) gives the number of active distinct colors and is appended to the result array.

The variable limit is not actually needed in the implementation because the queries are guaranteed to contain valid ball labels.

Go Solution

func queryResults(limit int, queries [][]int) []int {
	ballToColor := make(map[int]int)
	colorCount := make(map[int]int)

	result := make([]int, 0, len(queries))

	for _, query := range queries {
		ball := query[0]
		newColor := query[1]

		// Remove old color if ball already has one
		if oldColor, exists := ballToColor[ball]; exists {
			colorCount[oldColor]--

			if colorCount[oldColor] == 0 {
				delete(colorCount, oldColor)
			}
		}

		// Assign new color
		ballToColor[ball] = newColor

		// Increase new color frequency
		colorCount[newColor]++

		// Store distinct color count
		result = append(result, len(colorCount))
	}

	return result
}

The Go implementation mirrors the Python solution closely.

Go maps naturally support sparse storage, which is ideal because ball labels and colors can be very large.

One Go-specific detail is checking whether a ball already exists using:

if oldColor, exists := ballToColor[ball]; exists

This distinguishes between an absent key and a valid stored value.

Another detail is explicitly deleting colors whose frequency becomes zero using delete(colorCount, oldColor). This ensures that len(colorCount) accurately represents the number of distinct colors.

Worked Examples

Example 1

Input:

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

Initial state:

ball_to_color = {}
color_count = {}

Query	Action	ball_to_color	color_count	Distinct Colors
[1,4]	Assign color 4 to ball 1	{1:4}	{4:1}	1
[2,5]	Assign color 5 to ball 2	{1:4, 2:5}	{4:1, 5:1}	2
[1,3]	Remove color 4, assign 3	{1:3, 2:5}	{3:1, 5:1}	2
[3,4]	Assign color 4 to ball 3	{1:3, 2:5, 3:4}	{3:1, 5:1, 4:1}	3

Final output:

[1,2,2,3]

Example 2

Input:

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

Query	Action	ball_to_color	color_count	Distinct Colors
[0,1]	Assign color 1	{0:1}	{1:1}	1
[1,2]	Assign color 2	{0:1, 1:2}	{1:1, 2:1}	2
[2,2]	Assign color 2	{0:1, 1:2, 2:2}	{1:1, 2:2}	2
[3,4]	Assign color 4	{0:1, 1:2, 2:2, 3:4}	{1:1, 2:2, 4:1}	3
[4,5]	Assign color 5	{0:1, 1:2, 2:2, 3:4, 4:5}	{1:1, 2:2, 4:1, 5:1}	4

Final output:

[1,2,2,3,4]

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Each query performs only constant-time hash map operations
Space	O(n)	At most n balls and n colors are stored

Each query performs a fixed number of dictionary or hash map operations:

• Lookup
• Insert
• Delete
• Increment or decrement

These operations are all O(1) on average.

Since we process n queries exactly once, the total runtime is O(n).

The space complexity is also O(n) because, in the worst case, every query introduces a new ball and a new color.

Test Cases

from typing import List
from collections import defaultdict

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        ball_to_color = {}
        color_count = defaultdict(int)

        result = []

        for ball, new_color in queries:
            if ball in ball_to_color:
                old_color = ball_to_color[ball]

                color_count[old_color] -= 1

                if color_count[old_color] == 0:
                    del color_count[old_color]

            ball_to_color[ball] = new_color
            color_count[new_color] += 1

            result.append(len(color_count))

        return result

sol = Solution()

# Provided example 1
assert sol.queryResults(4, [[1,4],[2,5],[1,3],[3,4]]) == [1,2,2,3]

# Provided example 2
assert sol.queryResults(4, [[0,1],[1,2],[2,2],[3,4],[4,5]]) == [1,2,2,3,4]

# Recoloring removes old distinct color
assert sol.queryResults(2, [[0,1],[0,2]]) == [1,1]

# Multiple balls sharing same color
assert sol.queryResults(3, [[0,5],[1,5],[2,5]]) == [1,1,1]

# Repainting with same color
assert sol.queryResults(1, [[0,7],[0,7]]) == [1,1]

# Old color still exists after recolor
assert sol.queryResults(3, [[0,1],[1,1],[0,2]]) == [1,1,2]

# Single query
assert sol.queryResults(10, [[5,100]]) == [1]

# Large labels and colors
assert sol.queryResults(
    10**9,
    [[999999999, 1000000000], [888888888, 1000000000]]
) == [1,1]

# Color disappears completely
assert sol.queryResults(3, [[0,1],[1,2],[0,2]]) == [1,2,1]

# Many recolors
assert sol.queryResults(2, [[0,1],[0,2],[0,3],[0,4]]) == [1,1,1,1]

Test	Why
Example 1	Validates standard recoloring behavior
Example 2	Validates multiple balls sharing colors
Recoloring removes old color	Ensures disappearing colors are removed correctly
Multiple balls same color	Ensures duplicates are not double-counted
Repainting with same color	Checks idempotent updates
Old color still exists	Ensures frequencies are tracked correctly
Single query	Smallest meaningful input
Large labels and colors	Confirms hash-map based handling
Color disappears completely	Verifies deletion from frequency map
Many recolors	Tests repeated updates on same ball

Edge Cases

One important edge case occurs when a ball is recolored and its previous color disappears completely from the system. For example, if only one ball has color 4 and that ball changes to color 7, then color 4 must no longer count as distinct. A common bug is forgetting to remove colors whose frequency becomes zero. This implementation handles the issue by deleting the color from color_count whenever its count reaches zero.

Another important edge case is when multiple balls share the same color. If ball 0 and ball 1 both use color 5, the distinct color count should still only increase once. The frequency map naturally handles this because multiple balls increment the same counter instead of creating duplicate entries.

A subtle edge case occurs when a ball is recolored with the exact same color it already has. For example, changing ball 2 from color 8 to color 8 should not change the number of distinct colors. The implementation still works correctly because it decrements the old frequency and then increments it again, leaving the final count unchanged.

Another potential issue comes from the extremely large values allowed for ball labels and colors. Since values can be as large as 10^9, using arrays indexed by these values would be impossible. The implementation avoids this entirely by using hash maps, which store only the labels and colors that actually appear in queries.