백준 2399 - 거리의 합

www.acmicpc.net/problem/2399

 

C 계열 언어를 쓰면 $O(N^2)$로 뚝딱 풀 수 있지만 파이썬으로 풀어보았다.

풀이

점과 점 사이 거리를 누적합으로 저장하면 $O(N)$으로 모든 쌍의 거리를 저장할 수 있다.

 

점들의 좌표들이 정렬된 배열 ar에서 $f(i - 1)$이 ar[i - 1]과 이전 점들의 거리 합을 나타낼 때 ar[i]와 이전 점들의 거리 합을 구하는 점화식은 $f(i) = f(i - 1) + (ar_{i} + ar_{i - 1}) \times i$이다.

 

i가 1씩 증가할 때마다 이전 점들과 ar[i] 사이의 거리를 구하기 위해 $(ar_{i} + ar_{i - 1})$를 지금까지 거쳤던 점들의 개수만큼 더해줘야 한다는 사실을 생각하면 $(ar_{i} + ar_{i - 1}) \times i$의 의미를 쉽게 이해할 수 있다.

 

누적합으로 이를 저장하고 loop를 돌려 누적합의 값을 모두 더한다. 문제에서는 $N^2$ 쌍의 거리 합을 구하라고 했으므로 더해준 값에 2를 곱한 값이 정답이다.

 

시간 복잡도는 정렬을 해야 하므로 $O(N \log N)$이다.

 

전체 코드 - 파이썬

import sys
 
n = int(sys.stdin.readline().rstrip())
ar = sys.stdin.readline().rstrip().split(' ')
prefixsum = []
 
for i in range(n):
    ar[i] = int(ar[i])
    prefixsum.append(0)
 
ar.sort()
 
for i in range(1, n):
    prefixsum[i] = prefixsum[i - 1] + (ar[i] - ar[i - 1]) * i
 
res = 0
for i in prefixsum:
    res += i
 
print(res * 2)