10800. 컬러볼 (Python)

10800번: 컬러볼

 

풀이

 두 개의 정보에 대해서 저장하는 리스트를 만들어둔다. 하나는 색깔 별 무게를 담아두는 리스트이고 다른 하나는 같은 사이즈를 가진 공들의 정보를 담아놓는 리스트이다. 입력받은 공들에 대해서 사이즈 기준 오름차순으로 정렬한다. 그 후 차례대로 공들의 무게를 누적하여 더해나간다. 

 

 모든 공들의 무게를 더한 값에서 현재 공과 같은 색을 가진 공들의 무게를 빼주면 현재 공이 잡을 수 있는 공들의 크기의 합을 구할 수 있다. 같은 사이즈의 경우는 따로 리스트에 담아뒀다가 현재 공과 이전 공들의 무게가 달라졌을 경우 해당 리스트에 있는 공들을 전체 무게와 색깔 별 무게에 더해 주는 과정을 거치면 된다.

 

코드

import sys
import heapq
input = sys.stdin.readline

if __name__ == '__main__':
    N = int(input())
    ans = [0 for _ in range(N)] # 각 공이 잡을 수 있는 공들의 크기의 합을 담아두는 리스트
    colors = [0 for _ in range(N+1)] # 각 색깔별 공들의 크기의 합 

    q = []  # (크기, 색, 인덱스 번호)
    for i in range(N):
        C, S = map(int, input().split())
        heapq.heappush(q, (S, C, i))

    sizeSum = 0 # 모든 공들의 크기의 누적합
    sameSize = [] # 같은 크기를 가진 공들의 정보를 담아두는 리스트
    prevSize = 0 # sameSize 리스트에 담긴 공들의 크기

    for i in range(N):
        S, C, idx = heapq.heappop(q)

        # 현재 공의 사이즈와 prevSize가 다를 경우
        # sameSize에 담긴 공들을 꺼내어 sizeSum과 colors 리스트에 크기를 더해준다.
        if prevSize != S: 
            while sameSize:
                size, color = sameSize.pop()
                sizeSum += size
                colors[color] += size
        prevSize = S # prevSize를 현재 공의 사이즈로 변경 
        sameSize.append((S, C)) # 일단 sameSize 리스트에 담아둔다.
        ans[idx] = sizeSum - colors[C] # 모든 공들의 크기의 합 - 현재 공과 같은 색을 가진 공들의 크기의 합

    for i in range(N):
        print(ans[i])