ABC 334 복기 - 배열에서 원소 하나를 제거해서 두 원소 절댓값 차이의 합을 최소로 만들기(prefix sum, suffix sum)

1. 문제

 

C - Socks 2 (atcoder.jp)

C - Socks 2

 

2. 풀이

 

N쌍의 양말이 있는데, 이 중 지정된 K개의 양말은 각각 1개를 잃어버린 상태이다.

 

양말의 색을 1,2,3,...,N으로 표시할때, 현재 상태에서 적절하게 짝을 짓는다.

 

짝지은 양말이 i번째 색, j번째 색일때 |i-j| 의 value를 가진다.

 

모든 짝의 value합이 최소가 되도록 만든다면?

 

양말의 개수 2N-K가 짝수라면 상관없지만 홀수라면 1개를 제외하고 짝을 지어야한다.

 

직관적으로 2N-K가 짝수라면, 양말의 색이 오름차순이 되도록 만들고 서로 인접한 두 원소끼리 짝지으면 최적이다.

 

굳이 생각해보자면.. 양말 색은 오름차순으로 만든다면

 

1,1,2,2,3,3,4,4,...,N,N인데 여기서 지정된 A1 < A2 < A3 < ... < AK를 잃어버린다면..

 

1,1,2,2,3,3,4,4,...,N,N에서 N-K쌍은 짝이 맞는 잃어버리지 않은 양말끼리 있고,

 

나머지 K개는 A1,A2,..,AK가 된다.

 

짝이 맞는 잃어버리지 않은 양말끼리는 그 색깔 번호 차이가 0이고 이것이 최소이다.

 

K개는 A1 < A2 < A3 < ... < AK에 대하여... (A1,A2), (A3,A4), ... , (AK-1,AK)로 짝지어야 절댓값 차이 합이 최소가 된다.

 

왜냐하면 어떤 A1 < A2 < Ai < Aj인 i,j에 대하여...

 

|Ai - A1| + |Aj - A2| > |A2 - A1| + |Aj - Ai|이기 때문에 (Ai,A1), (Aj,A2)보다 (Ai,Aj), (A1,A2)로 짝짓는게 유리하다.

 

 

근데 이제 2N-K가 홀수여서 1개를 제외해야할때가 문제다.

 

어떤 양말을 제외해야 차이 합을 최소로 만드는가?

 

생각해볼 수 있는 방법은 브루트 포스로 모든 양말을 하나씩 제외해보고 차이 합을 계산해보는건데...

 

N이 최대 20만인데 최악의 경우 $O(N^{2})$이라 문제다.

 

O(N)에 해결할 수 있는 놀라운 방법은 다음과 같다

 

 

 

맨 마지막을 제거하는 경우 A,B,C를 계산하고

 

맨 처음을 제거하는 경우 X,Y,Z를 계산해둔다면... 놀랍게도 그 중간 과정들은 위와 같이,

 

먼저 맨 마지막을 제거하는 경우 A+B+C에서

 

C를 빼고, Z를 더한 A+B+Z

 

B를 빼고 Y를 더한 A+Y+Z

 

A를 빼고 X를 더한 X+Y+Z이므로, A,B,C,X,Y,Z로 O(N)에 만들 수 있다는 것이다.

 

#배열에서 한 원소를 제거하고 인접한 원소간 절댓값 차이 합을 최소로 만드는 방법
n,k = map(int,input().split())

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

#직관적으로 원소를 오름차순 정렬한 다음, 인접한 원소끼리 차이 합을 구하는게 최적이다.
#[1,1,2,2,3,3,...,N,N] 형태로 배열을 재구성
#A에 존재하는 K개의 원소는 제거해야하므로..
visited = [0]*(n+1)

for i in range(k):
    
    visited[A[i]] = 1 #A의 원소는 미리 표시해두고

B = []

for i in range(1,n+1):
    
    if visited[i] == 0: #A의 원소가 아니라면 한번 더 넣어주고

        B.append(i)

    B.append(i)    

answer = 0

#2N-K가 짝수라면, 그냥 B배열에서 i+1번째 원소 - i번째 원소 합이 최소이다.
if (2*n-k) % 2 == 0:

    for i in range(0,len(B)-1,2):

        answer += (B[i+1]-B[i])

#2n-k가 홀수라면...
else:
    
    s = []
    e = []
    
    #마지막 원소를 제외하고 i+1번째 - i번째 원소 차이를 모두 구해놓는다
    #s1,s2,s3,...,sn
    for i in range(0,len(B)-1,2):
        
        s.append(B[i+1]-B[i])
    
    answer = sum(s)
    
    #첫번째 원소를 제외하고 i+1번째 - i번째 원소 차이를 모두 구해놓는다
    #e1,e2,e3,...,en
    for i in range(1,len(B),2):
        
        e.append(B[i+1]-B[i])
    
    k = answer
    
    #s1+s2+s3+...+sn에서 시작해서, 
    #s1+s2+s3+...sn-1 +sn - sn + en
    #s1+s2+s3+...+sn-2+sn-1 - sn-1 + en-1 + en
    #s1+s2+s3+...+sn-3+sn-2 - sn-2 + en-2 + en-1 + en
    #... 으로 모두 조사해보면 된다
    for i in range(len(s)-1,-1,-1):

        k = k - s[i] + e[i]    
        
        if answer > k:
            
            answer = k
    
print(answer)