[python] 백준 10800 컬러볼

https://www.acmicpc.net/problem/10800

10800번: 컬러볼

 

문제 해결

• 우선 기준크기보다 크면 계산할 필요가 없으므로 sort()를 이용하여 공에 대한 정보를 공의 크기 오름차순으로 나타낸다.
• 크기가 기준크기보다 작아도 색이 같으면 계산에서 제외되므로 색깔별로 크기합에 대한 고려를 해야한다.
• 위의 생각을 통해 sort()를 한 후 for문으로 하나씩 기준을 잡는다. 이때 공의 크기가 작은 것부터 기준을 잡기 때문에 인덱스i 보다 작은 j에 대해서만 확인을 하면 된다.
• i보다 작은 j는 i+1공보다도 작으므로 total을 이용하여 계속 공의 합을 계산해 나간다. 그리고 dictionary를 통해 현재 인덱스i까지 보았을 때 색별로 공의 합을 value로 저장한다.
• 기준 i공에 대한 계산이 끝났을 때 마지막으로 출력할 값을 total-색이 같아서 공의 크기 합에 제외되는 공의 크기 합을 저장해 놓는다.

 

CODE

from collections import defaultdict
import sys
input = sys.stdin.readline
n = int(input())
A = []
for i in range(n):
    c, s = map(int, input().split())
    A.append([c,s,i])
A.sort(key=lambda x: x[1]) # 크기에 따라 정렬
answer = defaultdict(int)
sum_ball_size = defaultdict(int)

total = 0
j =0
for i in range(n):
    while A[j][1] < A[i][1]:
        total += A[j][1]
        sum_ball_size[A[j][0]] += A[j][1]
        j += 1
    answer[A[i][2]] = total - sum_ball_size[A[i][0]]
for i in range(n):
    print(answer[i])