Find the Number of Distinct Colors Among the Balls

main.md

Question

Approach-ish(es)

In this approach, I used a dictionary ball_colors to keep track of the latest color assigned to each ball. Also, I added a set called distinct_colors to store all unique colors in use.

For each query [x, y], I first check if ball x has already been assigned a color. If it has and the color is different from y, I decrement the count of the old color in a dictionary color_count. If the count reaches zero, I remove the old color from distinct_colors. Then, I update ball x with the new color y and update its count in color_count. Finally, I add y to distinct_colors and append the size of distinct_colors to my result list.

Implementation

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        ball_colors = {}
        distinct_colors = set()
        color_count = {}
        result = []

        for x, y in queries:
            if x in ball_colors:
                old_color = ball_colors[x]
                if old_color != y:
                    color_count[old_color] -= 1
                    if color_count[old_color] == 0:
                        distinct_colors.remove(old_color)

            ball_colors[x] = y
            if y not in color_count:
                color_count[y] = 0
            color_count[y] += 1
            distinct_colors.add(y)

            result.append(len(distinct_colors))

        return result

Complexities

• Time: O(n)
• Space: O(n)

The code didn't pass all testcases so I gave up and checked the editorial and it worked!