백준 10800 - 컬러볼 (파이썬)

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

풀이

각 공은 색 c와 크기 s를 가짐, 공을 하나 선택했을 때, 자신보다 크기가 작은 모든 공들 중 색이 다른 공들의 크기 합을 구해야 한다.

최대 크기인 2000을 기준으로 balls[크기] = [(색, 인덱스)] 형태로 공을 저장한다.

N = int(input())
balls = [[] for _ in range(2001)]

입력받은 공의 색(c)과 크기(s)를 balls 배열에 저장한다. (같은 크기의 공들을 묶어서 저장)

for i in range(N):
    c, s = map(int, input().split())
    balls[s].append((c, i))

누적합 계산을 진행한다. 단 작은 크기부터 순회한다.

for i in range(1, 2001):
    for c, id in balls[i]:
        if c not in count:
            count[c] = 0
        answer[id] = total - count[c]

현재 크기 i의 공들을 확인하며, 전체 누적 크기(total) 에서 같은 색 공들의 누적 크기(count[c]) 를 빼준다.

즉, 같은 색깔 공은 제외하고 자신보다 작은 크기의 공의 크기합을 구하는것이다.

현재 크기의 공들을 누적합에 추가한다.

total += i * len(balls[i])
for c, _ in balls[i]:
    count[c] += i

크기 i인 공이 n개라면 total += i * n 을 통해서 각 색깔별로도 누적합을 관리한다.

코드

# 백준 10800 - 컬러볼
# 분류 : 누적합

import sys
input = sys.stdin.readline

N = int(input())

balls = [[] for _ in range(2001)]

for i in range(N) :
    c, s = map(int, input().split())
    balls[s].append((c, i))

answer = [0] * N
total = 0
count = {}
for i in range(1, 2001) :
    for c, id in balls[i] :
        if c not in count :
            count[c] = 0
        answer[id] = total - count[c]
    
    total += i * len(balls[i])

    for c, _ in balls[i] :
        count[c] += i

for i in range(N) :
    print(answer[i])