[백준 10800] 컬러볼 - Python

문제 보기

 

[사용한 알고리즘]

정렬, 구현

 

[알고리즘]

1. 공을 크기 순으로 오름차순 정렬합니다.

2. 0번 인덱스 공부터 순차적으로 탐색합니다.

- 현재 인덱스까지 공 무게의 총합을 total 변수에 저장합니다.

- 각 색깔별로 공 무게 누적합을 ball_size_sum 딕셔너리에 저장합니다.

- 현재 인덱스까지 공 무게 총합 - 현재 색깔 공 무게 누적합은 현재 인덱스 플레이어가 사로잡을 수 있는 모든 공들의 크기의 합입니다.

 

※ 단순한 2중 for문으로 구현하면 n이 200,000으로 큰 수이기 때문에 시간 초과가 발생합니다. 따라서 안쪽의 for문의 횟수를 줄이기 위해 바깥쪽 for문의 인덱스를 순차적으로 증가시키면서 안쪽의 for문은 while문으로 변경하여 불필요한 반복은 수행하지 않도록 구현하였습니다. 

 

 

[코드]

from collections import defaultdict
import sys

if __name__ == "__main__":
    n = int(sys.stdin.readline())
    ball = []
    for i in range(n):
        c, s = map(int, sys.stdin.readline().split())
        ball.append([c, s, i])

    # 공의 크기 순으로 정렬
    ball.sort(key=lambda x: x[1])

    answer = defaultdict(int)
    ball_size_sum = defaultdict(int)  # 색깔 별 공의 크기 누적합

    total = 0  # 총합
    i, j = 0, 0
    for i in range(n):
        while ball[j][1] < ball[i][1]:  # 크기가 작을 때까지 수행
            total += ball[j][1]
            ball_size_sum[ball[j][0]] += ball[j][1]
            j += 1
        answer[ball[i][2]] = total - ball_size_sum[ball[i][0]]  # 총합 - 현재 색깔 공 누적합

    for i in range(n):
        print(answer[i])