[BOJ] 백준 10800번 컬러볼 - 파이썬(Python)

문제

지훈이가 최근에 즐기는 컴퓨터 게임이 있다. 이 게임은 여러 플레이어가 참여하며, 각 플레이어는 특정한 색과 크기를 가진 자기 공 하나를 조종하여 게임에 참여한다. 각 플레이어의 목표는 자기 공보다 크기가 작고 색이 다른 공을 사로잡아 그 공의 크기만큼의 점수를 얻는 것이다. 그리고 다른 공을 사로잡은 이후에도 본인의 공의 색과 크기는 변하지 않는다. 다음 예제는 네 개의 공이 있다. 편의상 색은 숫자로 표현한다.

공 번호	색	크기
1	1	10
2	3	15
3	1	3
4	4	8

이 경우, 2번 공은 다른 모든 공을 사로잡을 수 있다. 반면, 1번 공은 크기가 더 큰 2번 공과 색이 같은 3번 공은 잡을 수 없으며, 단지 4번 공만 잡을 수 있다. 

공들의 색과 크기가 주어졌을 때, 각 플레이어가 사로잡을 수 있는 모든 공들의 크기의 합을 출력하는 프로그램을 작성하시오. 

입력

첫 줄에는 공의 개수를 나타내는 자연수 N이 주어진다(1 ≤ N ≤ 200,000). 다음 N개의 줄 중 i번째 줄에는 i번째 공의 색을 나타내는 자연수 Ci와 그 크기를 나타내는 자연수 Si가 주어진다(1 ≤ Ci ≤ N, 1 ≤ Si ≤ 2,000). 서로 같은 크기 혹은 같은 색의 공들이 있을 수 있다.

출력

N개의 줄을 출력한다. N개의 줄 중 i번째 줄에는 i번째 공을 가진 플레이어가 잡을 수 있는 모든 공들의 크기 합을 출력한다.

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

10800번: 컬러볼

 

💡 풀이 및 코드

from collections import defaultdict
n = int(input())
data = []
for i in range(n):
    data.append(tuple(map(int, input().split()))+tuple([i]))
data.sort(key=lambda x: (x[1], x[0]))
sum_ball = 0
dict_ball = defaultdict(int)
dup_ball = defaultdict(int)
result, temp = [], (0, 0)
for color, size, i in data:
    if temp == (color, size):
        result.append((i, result[-1][1]))
    else:
        result.append((i, sum_ball-dict_ball[color]-dup_ball[size]))
    sum_ball += size
    dict_ball[color] += size
    dup_ball[size] += size
    temp = (color, size)
result.sort()
for i, r in result:
    print(r)

먼저, 크기와 색 값으로 오름차순 정렬한 후 각 원소의 인덱스를 idx라고할 때, 다음과 같이 사로잡을 수 있는 공 크기를 구한다.

idx이전의 모든 공 크기 합 - idx의 공과 색깔이 같은 공 크기 합 - idx의 공과 크기가 같은 공 크기 합

단, idx 공과  idx-1 공의 크기와 색이 같을 때는 위의 공식으로 계산하지 않고 idx 공이 idx-1 공의 결과 값을 그대로 가져갈 수 있도록 예외처리를 해주어야 한다.