[ BOJ / Python ] 10800번 컬러볼

황승환·2022년 7월 7일
0

Python

목록 보기
349/498


이번 문제는 누적합을 활용하여 해결하였다. 처음에는 단순하게 공 크기의 오름차순으로 정렬을 하였고, 2중 for문을 이용하여 매번 값을 계산하도록 하였다. 이럴 때 최악의 경우 200,000 x 200,000의 연산이 필요하기 때문에 시간 초과가 발생할 것을 알고는 있었다. 이제 시간을 줄이는 방법에 대해 생각하다가 값을 누적하면 되겠다고 생각하였다. 여기서 문제는 같은 색의 공일 경우에는 값을 빼야 하는데 이를 어떻게 관리하느냐였다.

고민 끝에 색깔 별로 누적합을 같이 구한다면 가능할 것 같다는 생각을 하였다. 그래서 정렬된 공 리스트를 순회하며 현재 공의 값을 누적합에 더해주고, 이때 현재 공의 색에 해당하는 리스트에도 값을 더해주었다. 그리고 결과 리스트에는 현재까지의 누적합 - 현재 공 색에 해당하는 누적합의 값을 넣어주었다.

Code

n = int(input())
balls = []
for i in range(n):
    a, b = map(int, input().split())
    balls.append((i, a, b))
balls.sort(key = lambda x: (x[2], x[1]))
totals = [0 for _ in range(200001)]
answer = [0 for _ in range(n)]
total, i, j = 0, 0, 0
while True:
    if i >= n:
        break
    b1, b2 = balls[i], balls[j]
    while b2[2] < b1[2]:
        total += b2[2]
        totals[b2[1]] += b2[2]
        j += 1
        b2 = balls[j]
    answer[b1[0]] = total - totals[b1[1]]
    i += 1
for i in range(n):
    print(answer[i])

profile
꾸준함을 꿈꾸는 SW 전공 학부생의 개발 일기

0개의 댓글