[Python] 백준 10800번: 컬러볼

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

10800번: 컬러볼

import sys
input = sys.stdin.readline


def getSumsToPrint(inputs):
    sortedInputs = sorted(
        sorted(inputs, key=lambda x: x[0]), key=lambda x: x[1])
    prefixSums = [[0] for _ in range(N+1)]

    prevSize = 0
    prevColor = 0
    prevZeroIndex = 0

    sumsToPrint = {}

    for i in range(N):
        c, s = sortedInputs[i]
        prefixSums[0].append(prefixSums[0][-1] + s)
        prefixSums[c].append(prefixSums[c][-1] + s)

        if s == prevSize:
            if c != prevColor:
                sumsToPrint[c, s] = prefixSums[0][prevZeroIndex] - \
                    prefixSums[c][-1]
        else:
            sumsToPrint[c, s] = prefixSums[0][-1] - prefixSums[c][-1]
            prevZeroIndex = i+1

        prevSize = s
        prevColor = c

    return sumsToPrint


if(__name__ == '__main__'):
    N = int(input())

    # 입력
    inputs = []
    for _ in range(N):
        inputs.append(tuple(map(int, input().split())))

    # 처리
    sumsToPrint = getSumsToPrint(inputs)

    # 출력
    for input in inputs:
        print(sumsToPrint[input])

상당히 어려운 문제이다. 수많은 접근법을 고안해낸 끝에 올바른 접근방식을 찾아냈다. 바로, 크기순으로 공들을 정렬시킨 뒤, 작은 것부터 시작해서 반복문을 돌면서, 공 전체의 누적합과 각 색상의 누적합을 따로 저장한다는 것이다. 이 때 특정 공에 대해서 누적합을 저장한 시점에서, 전체의 누적합에서 해당 공의 색의 누적합을 빼주면 그것이 우리가 구해주는 값이 된다. 모든 공들의 누적합에서 해당 색상에서의 누적합을 뺀다는 것이 무엇을 의미하는지 생각해보라.

 

이것을 목표로 하고 코드를 짜면 된다. getSumsToPrint()를 통해 sumsToPrint라는 딕셔너리를 반환받고 이것을 출력하는 것이 큰 틀이다.

 

getSumsToPrint() 내부를 설명하겠다. 우선 정렬을 크기에 대해서 1번이 아니라, 색상에 대해서도 정렬을 또 해주었는데, 이것이 필요한 이유는 나중에 나온다.

for문 내부를 살펴보자. 우선 prefixSums[0]은 전체 공들의 누적합이고, prefixSums[n]은 n의 색에 대한 누적합을 의미한다. 따라서, 2~3번째 줄은 각각 전체 공의 누적합과 특정 색상의 누적합을 추가해주는 의미이다.

다음은 현재 크기(s)와 이전 크기(prevSize)를 비교하며 분기한다. 우선 사이즈가 다를 때부터 보면, 가장 일반적인 코드가 이곳에 해당되는데, 앞서 설명한 바와 같이, 그저 지금까지 저장된 prefixSums를 기준으로 전체 공의 누적합에서 해당 공의 누적합을 빼주고, prevZeroIndex라는 것을 증가시켜줄 뿐이다. 이후엔 분기와는 별개로 prevSize와 prevColor라는 것을 갱신시킨다.

그럼, prevSize와 s가 같은 경우를 보자. 이 안에서는 또 현재 색상(c)과 이전 색상을 비교하는 분기점이 있다. 우선 크기와 색상이 모두 같을 경우에는 아무런 동작도 하지 않는다. 여기서 sortedInputs가 정렬을 색상에 대해서까지 굳이 2번 한 이유가 설명되는데, 가령 공이 2 5 / 1 10 / 1 10이 들어왔다고 하자. 이 때 색이 1인 두 공은 모두 5라는, 같은 값을 가지기 때문이다.

다음은 색이 같으나 색상이 다를 경우인데, 이 때는 특이하게도 prefixSums[0][-1]이 아니라 [0][prevZeroIndex]가 들어온다. 1 5 / 3 5 / 2 10 / 3 10이 들어왔다고 하면, 맨 뒤의 2개의 공은 각각, 10과 5가 들어올 것이다. 그러니까, (모든 공의 누적합) - (현재 색의 누적합)에서 모든 공의 누적합 부분은 중복되는 크기가 들어온 처음 상태일 때의 값으로 들어간다는 것이다.