ABC 351 F번 복기 - 알고리즘 문제에 max함수를 바꾸는 트릭 2가지

F - Double Sum (atcoder.jp)

F - Double Sum

 

문제는 매우 간단하다

 

A1,A2,A3,...,AN이 주어지면, $$\sum_{i = 1}^{n} \sum_{j = i+1}^{n} max(A_{j} - A_{i},0)$$을 구하는 문제

 

n제한이 40만이라 단순하게 풀면 당연히 시간초과...

 

 

1. max(a,b) = (|a-b| + a+b)/2

 

방법은 많이 있던데 아주 간단하고 경이로운 솔루션이 있어서 복기해본다

 

배열 A에 대한 함수 f를 다음과 같이 정의한다.

 

$$f(A) = \sum_{i = 1}^{n} \sum_{j = i+1}^{n} (A_{j} - A_{i})$$

 

식을 정리하면

 

$$f(A) = \sum_{i = 1}^{n} ( \sum_{j = i+1}^{n} A_{j} - \sum_{j = i+1}^{n} A_{i} )$$

 

$A_{i}$는 j에 depend하지 않으므로

 

$$f(A) = \sum_{i = 1}^{n} (\sum_{j = i+1}^{n} A_{j} - (n - i)A_{i})$$

 

$\sum_{j = i+1}^{n} A_{j} = prefix[n] - prefix[i]$와 같다.

 

여기서 prefix[i] = A[1] + A[2] + ... + A[i] 누적합

 

따라서

 

$$f(A) = \sum_{i = 1}^{n} (prefix[n] - prefix[i] - (n-i)A_{i})$$

 

그러므로 prefix를 먼저 구해놓는다면, 함수 f(A)는 O(N)에 구할 수 있다.

 

 

여기서

 

$$max(A_{j} - A_{i}, 0) = \frac{|A_{j} - A_{i}| + A_{j} - A_{i}}{2}$$

 

따라서 다음 식을 구해서 2로 나누면 된다.

 

$$\sum_{i = 1}^{n} \sum_{j = i+1}^{n} |A_{j} - A_{i}| + \sum_{i = 1}^{n} \sum_{j = i+1}^{n} (A_{j} - A_{i})$$

 

우항은 f(A)와 같은데 좌항은?

 

배열 A내에서 서로 다른 두 항 Aj와 Ai의 차이를 양수로 바꾸고 누적한 합인데,

 

Ai - Aj나 Aj - Ai나 절댓값은 같고 부호만 다르므로,

 

A를 오름차순 정렬하면 뒤쪽에 있는 값에서 앞쪽에 있는 값을 빼면 항상 양수가 되므로, A를 정렬할 때 A'이라고 한다면,

 

f(A')과 같다.

 

 

def f(A):
    
    n = len(A)

    prefix = [0]

    for i in range(n):
        
        prefix.append(prefix[-1] + A[i])
    
    v = 0

    for i in range(n):
        
        v += (prefix[n] - prefix[i] - (n-i)*A[i])
    
    return v

n = int(input())

A = list(map(int,input().split()))

B = sorted(A)

print((f(B) + f(A))//2)

 

 

 

2. max(a-b,0) = max(a,b) - b

 

더 간단한 방법이 있다는 것 같은데

 

$$max(A_{j} - A_{i},0) = max(A_{j},A_{i}) - A_{i}$$이다.

 

 

그러므로, 구하고자 하는 합은

 

$$\sum_{i=1}^{n} \sum_{j = i+1}^{n} max(A_{j},A_{i}) - A_{i} = \sum_{i=1}^{n} \sum_{j=i+1}^{n} max(A_{j},A_{i}) - \sum_{i=1}^{n} (n-i)A_{i}$$

 

여기서 좌항 $\sum_{i=1}^{n} \sum_{j=i+1}^{n} max(A_{j},A_{i})$는, 모든 i,j에 대하여 Aj,Ai중 최댓값을 의미하므로,

 

A를 정렬해도 그 값은 바뀌지 않는다.

 

 

따라서, A를 오름차순 정렬하면, 항상 Aj가 더 크므로,

 

$$\sum_{i = 1}^{n} \sum_{j=i+1}^{n} max(A_{j},A_{i}) = \sum_{i=1}^{n} \sum_{j=i+1}^{n} B_{j}$$

 

 

여기서 prefix를 유도해서 $\sum_{j = i+1}^{n} B_{j} = prefix[n] - prefix[i]$해도 되긴 하는데..

 

 

i = 1부터 n까지, j = i+1부터 n까지는....

 

i = 1, j = 2,3,4,..,n

 

i = 2, j = 3,4,...,n

 

i = 3,j = 4,5,...,n

 

...

 

j = 2부터 n까지, i = 1부터 j-1까지와 같다.

 

j = 2, i = 1

 

j = 3, i = 1,2

 

....

 

따라서 i,j를 뒤바꾸면 Bj는 i에 depend하지 않으므로...

 

$$\sum_{j = 2}^{n} \sum_{i = 1}^{j-1} B_{j} = \sum_{j = 2}^{n} (j-1)B_{j}$$

 

이 값은 A를 정렬하면 O(NlogN)에 찾을 수 있다.

 

 

n = int(input())

A = [0] + list(map(int,input().split()))
B = sorted(A)

v1 = 0

for j in range(2,n+1):
    
    v1 += (j-1 * B[j])

v2 = 0

for i in range(1,n+1):
    
    v2 += ((n-i) * A[i])

print(v1-v2)