diff --git a/medium/find_the_number_of_distinct_colors_among_the_balls.py b/medium/find_the_number_of_distinct_colors_among_the_balls.py new file mode 100644 index 0000000..5f2890a --- /dev/null +++ b/medium/find_the_number_of_distinct_colors_among_the_balls.py @@ -0,0 +1,20 @@ +# https://leetcode.com/problems/find-the-number-of-distinct-colors-among-the-balls + +from typing import Dict, List + +class Solution: + def queryResults(self, _: int, queries: List[List[int]]) -> List[int]: + balls: Dict[int, int] = {} + colors: Dict[int, int] = {} + result = [] + for [ball, color] in queries: + curr_color = balls.get(ball, 0) + if curr_color != 0: + colors[curr_color] -= 1 + if colors[curr_color] == 0: + del(colors[curr_color]) + balls[ball] = color + colors[color] = colors.get(color, 0) + 1 + result.append(len(colors)) + + return result