[LeetCode] 3160. Find the Number of Distinct Colors Among the Balls (Python)

Problem

 

Solution

hash

balls: x위치에 어떤 color 인지 저장

count_colors: 특정 color 가 몇번 등장했는지 count

 

만약 x 위치의 공이 이미 색칠되어있다면, 해당 색상 - 1

 

Code

class Solution:
    def queryResults(self, limit: int, queries: List[List[int]]) -> List[int]:
        n = len(queries)
        balls = dict()
        count_colors = defaultdict(int)
        answer = []
        cnt = 0

        for i in range(n):
            x, color = queries[i]

            if x in balls:
                # if already colored x, changed color
                # count_color[color] - 1 
                count_colors[balls[x]] -= 1
                #특정 색상 수가 0 이되면 dict에서 제거
                if count_colors[balls[x]] == 0:
                    del count_colors[balls[x]]

			#x 위치에 색깔 업데이트
            balls[x] = color
            count_colors[color] += 1
            #서로 다른 색상 개수 추가
            answer.append(len(count_colors))
            
        return answer

 

Complexity

Time Complexity

O(N)

 

Space Complexity

balls , count_colors , answer -> n개 저장

O(N)