[acmicpc] 10800. 컬러볼(python)

- 구간합 알고리즘

N이 200,000까지라서 이중 for문으로 구할 경우, 100억이 넘어가서 당연히 시간초과가 될 것이다. 따라서 구간합 알고리즘을 사용하는데, 이는 시간복잡도가 O(n)이다.

추가테스트 케이스  > 색과 무게가 겹치는 여러 경우를 고려하여 구성했다.

6
2 3
1 10
2 10
1 10
3 11
4 14

정답

0
3
0
3
33
44

import sys
input = sys.stdin.readline

N = int(input())
arr = []
result = [0] * (N + 1)
color = [0] * (N + 1)
weight = [0] * 2001
for i in range(N):
    e1, e2 = map(int, input().split())
    arr.append((e2,e1,i))
arr.sort() # 정렬은 무게, 색깔순으로
prefix_sum = [0] * (N + 1)

for i in range(0, N):
    prefix_sum[i] = prefix_sum[i - 1] + arr[i][0]
    
    if i>0 and arr[i][1]==arr[i-1][1] and arr[i][0]==arr[i-1][0]:
        # 색과 무게가 동시에 같을 경우 이전 값으로
        result[arr[i][2]]=result[arr[i-1][2]]
    else:
    	# 같은 색 빼기, 같은 무게 빼기
        result[arr[i][2]]=prefix_sum[i-1]-color[arr[i][1]]-weight[arr[i][0]]
    
    color[arr[i][1]]+=arr[i][0]
    weight[arr[i][0]]+=arr[i][0]

for i in range(0, N):
    print(result[i])

 

---

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

10800번: 컬러볼