Problem of The Day: Find the Number of Distinct Colors Among the Balls

Problem Statement

Brute Force [MLE]

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        N = len(queries)
        arr = [None] * (limit + 1)
        res = []
        for x, y in queries:
            arr[x] = y
            res.append(len(set([val for val in arr if val is not None])))
        return res

Hash Map Approach [TLE]

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        N = len(queries)
        num_of_balls = limit + 1
        colors = defaultdict(int)
        res = []
        for x, y in queries:
            colors[x] = y
            res.append(len(set(colors.values())))

        return res

Improved Algorithm

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        colors = {}  # Mapping of box -> color
        color_freq = defaultdict(int)  # Frequency of each color
        unique_count = 0  # Count of unique colors
        res = []

        for x, y in queries:
            if x in colors:
                old_color = colors[x]
                color_freq[old_color] -= 1
                if color_freq[old_color] == 0:
                    unique_count -= 1  # If color count drops to 0, it's no longer unique

            # Assign new color
            colors[x] = y
            if color_freq[y] == 0:
                unique_count += 1  # If it's a new color, increase unique count
            color_freq[y] += 1

            res.append(unique_count)

        return res

Editorial

Approach 1: Hashmap and Array

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        n = len(queries)
        result = []
        color_map = {}
        ball_array = [0] * (limit + 1)

        # Iterate through queries
        for i in range(n):
            # Extract ball label and color from query
            ball, color = queries[i]

            # Check if ball is already colored
            if ball_array[ball] != 0:
                # Decrement count of the previous color on the ball
                prev_color = ball_array[ball]
                color_map[prev_color] -= 1

                # If there are no balls with previous color left, remove color from color map
                if color_map[prev_color] == 0:
                    del color_map[prev_color]

            # Set color of ball to the new color
            ball_array[ball] = color

            # Increment the count of the new color
            color_map[color] = color_map.get(color, 0) + 1

            result.append(len(color_map))

        return result

Approach 2: Two Hash Maps

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        n = len(queries)
        result = []
        color_map = {}
        ball_map = {}

        # Iterate through queries
        for i in range(n):
            # Extract ball label and color from query
            ball, color = queries[i]

            # Check if ball is already colored
            if ball in ball_map:
                # Decrement count of the previous color on the ball
                prev_color = ball_map[ball]
                color_map[prev_color] -= 1

                # If there are no balls with previous color left, remove color from color map
                if color_map[prev_color] == 0:
                    del color_map[prev_color]

            # Set color of ball to the new color
            ball_map[ball] = color

            # Increment the count of the new color
            color_map[color] = color_map.get(color, 0) + 1

            result.append(len(color_map))

        return result