[ BOJ ] 10800 : 컬러볼 ( GOLD 3 ) / 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번째 공을 가진 플레이어가 잡을 수 있는 모든 공들의 크기 합을 출력한다.

 

풀이 과정

다른 공들이 크기가 작고 색이 다른지 전부 검사하며 결과 더하는 방법은 공이 최대 200,000개이기에 검사하는 데에 n^2 번의 연산이 진행되기에 1초 시간 제한으로 인하여 시간 초과가 날 것이라고 판단했다.

 

공을 크기 순으로 정렬한 후, 특정한 공보다 크기가 작은 공들 중에서 색이 다른 공들만 더하는 방법도 생각해보다가, 특정한 공보다 크기가 작은 공들의 크기 합에서 색이 같은 공들 중 특정한 공보다 크기가 작은 공들의 크기 합을 빼는 방법을 생각해냈다.

 

좋은 아이디어라고 생각하고 무작정 만들고, 색이 같은 공들 중 특정한 공보다 크기가 작은 공을 찾는 과정은 색들끼리 데이터를 정리한 뒤 while 문으로 완전 탐색을 이용해 크기가 작은 공을 찾기로 했었지만, 데이터가 중복될 가능성도 고려해야 했었다.

 

반복문을 이용한 무작정 탐색은 채택하지 않고, 전체 데이터와 색깔별 데이터에서 공의 크기가 처음 시작되는 부분을 기록하여, 처음 시작되는 부분보다 한 단계 작은 곳 까지의 합을 구하는 방식으로 변경하였고, 전체 데이터에서 크기가 중복되는 경우에도 이 방식으로 변경하니 AC를 받을 수 있었다.

필기 풀이

import sys
input = sys.stdin.readline

n = int(input())

# (공 번호, 색깔, 크기)
datas = []

for i in range(n):
    color, size = map(int, input().split())
    datas.append((i + 1, color, size))

# 크기 순으로 정렬
datas.sort(key=lambda x:x[2])

all_sums = [0] # 전체 공 크기의 누적 합
data_start = {} # 공 크기의 처음 시작 부분

nums = {} # 색깔 별 공들의 데이터
sums = {} # 색깔 별 공 크기의 누적 합
nums_start = {} # 색깔 별 공 크기의 처음 시작 부분

for i, c, s in datas:
    # 색깔 별로 공 데이터 누적하기
    if (sums.get(c, -1) == -1):
        sums[c] = [0]
        nums[c] = []

    sums[c].append(sums[c][-1] + s)
    nums[c].append(s)
    
    # 색깔 별 중복된 크기를 방지하기 위해 색깔 별 공의 크기 시작 부분 기록
    if (nums_start.get(c, -1) == -1):
        nums_start[c] = {}

    if (nums_start[c].get(s, -1) == -1):
        nums_start[c][s] = len(nums[c]) - 1

    # 중복된 크기를 방지하기 위해 공의 크기 시작 부분 기록
    if (data_start.get(s, -1) == -1):
        data_start[s] = len(all_sums) - 1

    # 누적 합
    all_sums.append(all_sums[-1] + s)

result = []

for i in range(n):
    number, color, size = datas[i]

    over = data_start[size] 
    # 해당 공의 크기가 처음 시작된 부분임.
    # -> all_sums에서는 0번째가 0부터 시작하기에 해당 공보다 한 단계 작은 공의 누적 합을 가리킴.

    idx = nums_start[color][size]
    # 같은 색깔 공의 크기가 처음 시작된 부분임.
    # -> sums에서는 0번째가 0부터 시작하기에 해당 색깔의 공보다 한 단계 작은 공의 누적 합을 가리킴.

    # 해당 공보다 크기가 작은 공들의 크기 합 - 색이 같은 공들 중 해당 공보다 크기가 작은 공들의 크기 합
    result.append([datas[i][0], all_sums[over] - sums[color][idx]])

# 공 번호 순서대로 재정렬 및 출력
result.sort(key=lambda x:x[0])
for i in range(n): print(result[i][1])

아이디어를 떠올리는 데에 시간이 오래 걸리지는 않았던 문제였다. 파이썬의 lambda 를 이용한 key 정렬을 배울 수 있었다.