Category
알고리즘  구현

[10800] 컬러볼

by Hwang

문제

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

풀이

만약 색이 다른 공을 사로잡아야한다는 조건이 없다면 무게를 기준으로 오름차순 정렬한 뒤 값을 누적해가며 정답을 구해줄 수 있습니다. (볼의 무게 = 이전까지의 누적 값)

여기서 같은 색을 잡을 수 있는 조건이 추가된다면, 무게 별 누적 값도 같이 구해준다면 같은 색이 제외된 누적 값을 구할 수 있게됩니다. (볼의 무게 = 이전까지의 누적 값 - 이전까지의 같은 색깔의 볼의 누적 값)

여기서 하나 생각해야하는 부분은 이전까지의 누적된 값에서 같은 무게인 경우를 제외시켜줘야 합니다.. 예를 들어 (1, 4), (2, 4)와 같은 조건이 입력 되었다면 0, 0의 값이 출력되어야하지만 지금과 같은 방식으로는 0, 4가 출력이 됩니다. 공 (2, 4)를 처리해줄 때 같은 무게인 (1, 4)가 제외되지 않았기 때문입니다.

따라서 같은 무게일 경우엔 전체 누적 값과, 같은 색깔의 볼의 누적 값을 갱신하지 않는 방식으로 풀이했습니다.

import sys

N = int(sys.stdin.readline())
balls = []
ball_weights = [0] * (N + 1)  # 누적된 볼의 무게
total = 0

for i in range(N):
    num, weight = map(int, sys.stdin.readline().split())
    balls.append((weight, num, i))

sorted_balls = sorted(balls)
temp = [0, 0]  # 같은 무게를 제외시켜주기 위한 배열 [무게, 현재까지 누적 무게]
stack = []  # 같은 무게를 넣어 줄 스택 

for w, n, i in sorted_balls:
    if temp[0] != w:
        while stack:
            num = stack.pop()
            ball_weights[num] += temp[0]
        temp[0], temp[1] = 0, 0

    balls[i] = (total-temp[1]) - ball_weights[n]
    total += w
    stack.append(n)

    if temp[0] == 0:
        temp[0], temp[1] = w, w
    else:
        temp[1] += w

for w in balls:
    print(w)


 Comment