3160. Find the Number of Distinct Colors Among the Balls

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/find-the-number-of-distinct-colors-among-the-balls
• topics: Hash Table

Description

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.

Idea

• Use 2 default dicts

Code

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        color_of_balls = defaultdict(int)
        nb_balls_per_color = defaultdict(int)
        ans = []
 
        for ball, color in queries:
            if ball in color_of_balls:
                prev_color = color_of_balls[ball]
                if prev_color == color:
                    ans.append(len(nb_balls_per_color))
                    continue
                nb_balls_per_color[prev_color] -= 1
                if not nb_balls_per_color[prev_color]:
                    del nb_balls_per_color[prev_color]
                
            color_of_balls[ball] = color
            nb_balls_per_color[color] += 1
 
            ans.append(len(nb_balls_per_color))
        return ans