[BOJ] 10800 - 컬러볼 (Python)

 

 

아이디어

 

현재 공보다 크기가 작은 공들의 크기 합과 크기가 작으면서 색이 같은 공들의 크기 합을 구하여 풀 수 있다.

 

 

풀이

 

공을 저장하는 리스트 balls에 공의 크기, 공의 색, 인덱스를 순서대로 추가하고 공의 크기 오름차순으로 정렬한다.

차례로 현재 공보다 작은 공의 크기를 total에 누적하여 더하고, sum_color 리스트에 공의 색에 따라 누적하여 더한다.

작은 공들의 크기 합에서 작으면서 색이 같은 공들의 크기 합을 뺀 값을 인덱스에 맞게 결과값으로 저장한다.

 

import sys

input = sys.stdin.readline

N = int(input())
balls = []
for i in range(N):
    a, b = map(int, input().split())
    balls.append([b, a, i])
balls.sort(key=lambda x: x[0])
sum_color = [0]*(N+1)
res = [0]*N
idx = total = 0
for i in range(N):
    while balls[idx][0] < balls[i][0]:
        total += balls[idx][0]
        sum_color[balls[idx][1]] += balls[idx][0]
        idx += 1
    res[balls[i][2]] = total-sum_color[balls[i][1]]
for i in range(N):
    print(res[i])

 

 

입력 받은 값들을 언패킹하는 리스트 컴프리헨션보다 for문에서 append하는 코드가 메모리가 더 적게 사용되고 더 빨랐다.

# 1.
balls = [[i, *map(int, input().split())] for i in range(N)]

# 2. 메모리가 더 적게 사용되고 더 빠른 코드
balls = []
for i in range(N):
    a, b = map(int, input().split())
    balls.append([b, a, i])

 

정렬할 때 key에 람다식으로 정렬할 부분을 지정한 코드가 더 빨랐다.

# 1.
balls.sort()

# 2. 더 빠른 코드
balls.sort(key=lambda x: x[0])

 

리스트를 언패킹하고 개행 문자를 구분자로 하는 코드보다 for문으로 출력하는 코드의 메모리가 더 적게 사용되었다.

# 1.
print(*res, sep='\n')

# 2. 메모리가 더 적게 사용되는 코드
for i in range(N):
    print(res[i])