컬러볼 (10800번) - 백준 (BOJ)

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

10800번: 컬러볼

 

백준 - 컬러볼 (10800번)

 

난이도 : Gold 3

알고리즘 : Sorting (정렬) & Prefix Sum (누적합)

풀이 소요 시간 : 33분

Time Complexity : O(NlogN)

 

import sys
input = sys.stdin.readline


if __name__ == '__main__':
    N = int(input())
    stack = []
    size_by_color = [0] * 200_001
    prev = total = 0
    balls = [(*map(int, input().split()), i) for i in range(N)]
    balls.sort(key=lambda x: x[1])
    answer = [0] * N
    
    for color, size, i in balls:
        if size == prev:
            answer[i] = total - size_by_color[color]
            stack.append((color, size))
            continue
        
        for cc, ss in stack:
            total += ss
            size_by_color[cc] += ss
        
        stack = [(color, size)]
        answer[i] = total - size_by_color[color]
        prev = size
        
    print('\n'.join(map(str, answer)))

 

Sorting (정렬) 과 Prefix Sum (누적합) 을 적절히 활용하면 해결할 수 있는 문제였다.

 

처음 공의 번호를 기억하기 위해,

각 공의 정보를 (color, size, #number) 형식으로 만들어준다.

ex) (1, 10), (3, 15), (1, 3), (4, 8) →  (1, 10, 0), (3, 15, 1), (1, 3, 2), (4, 8, 3)

 

그 후 size에 대해 정렬하고, 작은 크기의 공부터 하나씩 누적해서 더해가면 된다.

 

color 별로 저장할 수 있는 list를 만들어서 (size_by_color),

color 별 누적 합 (size_by_color[color])과 전체 누적 합 (total) 을 계속해서 업데이트 한다.

 

같은 size가 연속해서 여러 번 나오는 경우를 처리해줘야 하는데,

list를 하나 따로 만들어서 (stack)

바로 전 단계와 size가 동일하다면 해당 list (stack)에 저장해두었다.

 

그리고 바로 전 단계와 다른 size가 나타난다면,

해당 list (stack)에 저장된 값들을 하나씩 꺼내서, (size_by_color)와 (total)을 업데이트 한다.

 

최종적으로, 현재 단계의 공이 잡을 수 있는 공들의 크기 총 합은

(total - size_by_color[현재 공의 color]) 가 되고,

이를 (answer[현재 공의 #number])에 저장한다.

 

이렇게 하면 answer에 각 공이 잡을 수 있는 공들의 크기가 담기게 된다.