[Python] 백준 1150 - 백업

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

 

문제

당신은 큰 회사들의 컴퓨터 자료를 백업하여주는 정보통신회사를 운영한다. 자료 백업이 즐거운 일이 아니므로, 당신은 서로 다른 두 회사의 자료를 서로 백업하여 주는 시스템을 개발하여 당신이 집에서 게임을 즐기는 동안 백업이 이루어지게 하려 한다.

모든 회사들은 직선인 길을 따라 위치하고 있다. 당신은 서로 백업을 하여 줄 두 회사를 짝 지어 주어야 하는데, 두 회사 사이에 네트워크 케이블을 연결 사용하여야 한다.

네트워크 케이블은 대단히 비쌀 뿐 아니라, 지역 통신 회사에서는 당신에게 오직 k개의 네트워크 케이블을 제공한다 –이 말은 당신이 오직 k 쌍의 회사에만 백업을 할 수 있다는 뜻이다 (전체 2k 개의 회사). 어떤 회사도 두 개 이상의 쌍에 속할 수는 없다 (즉, 여기 2k 개의 회사가 모두 다른 회사라는 것을 뜻한다).

통신 회사는 네트워크 케이블의 길이를 km 단위로 경비를 부과한다. 따라서 당신은 가장 짧은 길이의 케이블을 사용하도록 회사들을 k 쌍으로 짝지어야 한다. 다시 말하자면, 회사들을 짝짓기 하는데, 짝지어진 두 회사간의 거리들의 전체 합을 최소화 하도록 짝을 지어야 한다는 것이다.

예를 들어, 아래 그림과 같이 다섯 개의 고객회사들이 같은 길 위에 위치한다고 하자. 이 회사들은 각각 길의 출발점에서 기준하여 1 km, 3 km, 4 km, 6 km 그리고 12 km 에 위치하고 있다. 통신회사에서는 오직 k = 2 케이블만을 제공한다.

예에서 최선의 짝 짓기 방법은 첫 번째와 두 번째 회사, 그리고 세 번째와 네 번째 회사를 묶어 주는 것이다. k = 2 케이블 만을 사용하게 되며, 첫 번째 케이블의 길이는 3 km –1 km = 2 km 이고, 두 번째 케이블의 길이는 6 km –4 km = 2 km 이다. 이와 같은 짝짓기는 전체 4 km 의 네트워크 케이블을 사용하게 되며, 실제 가능한 가장 짧은 경우이다.

입력

입력의 첫 번째 줄에는 정수 n 과 k 가 주어지는데, 각각 길 위의 회사 수(2 ≤ n≤ 100 000), 그리고 제공되는 네트워크 케이블 수(1 ≤ k ≤ ½ n)를 의미한다. 그 다음 n 줄에는 각각 하나의 정수 s (0 ≤ s ≤ 1 000 000 000) 가 주어지며, 이 정수는 길의 출발점에서 각 회사까지의 거리를 의미한다. 이 정수들은 가장 작은 것에서 가장 큰 것 까지 순서대로 나타난다. 어떤 두 개의 회사도 같은 지점에 있지 않다.

출력

출력은 반드시 하나의 양의 정수로 표현되어야 하는데, 이것은 2k 의 서로 다른 회사를 k 쌍으로 묶었을 때 필요한 가장 짧은 전체 네트워크 케이블의 길이이다.

 

---

 

풀이

대표적인 그리디 알고리즘 문제이다.

알고리즘 스터디에서 예전에 본 이후로 '풀어야지' 하고 까먹고 있었다가 오늘 드디어 풀게 되었다.

 

이 문제를 어떻게 그리디로 풀 수 있는가?

 

일단 관찰(이라기에는 너무 당연하다)을 통해 다음과 같은 사실을 알 수 있다:

A-B-C-D가 순서대로 있을 때, A-B 와 C-D를 고르는 것이 A-D 와 B-C를 고르는 것 보다 확실히 낫다.

 

그리고 문제 조건상, 한 회사가 두 번 연결될 수 없으므로 A-B-C와 같이 세 개 이상의 회사가 일렬로 쭉 연결되는 경우는 없다.

 

그러면 우리는 다음과 같은 방법을 생각해볼 수 있다.

 

"그리디하게 가장 짧은 경로들 우선으로 K개 뽑아주면 안되나요?"

 

그러면 이 문제가 D4일리가 없지. 당연히 다음과 같은 반례가 생긴다.

1-2를 뽑으면 0-1과 2-3은 선택 불능 상태가 되기 때문에 어쩔 수 없이 3-4를 뽑게 되지만, 이는 최적해가 아니며 0-1과 2-3을 뽑는 것이 최적해가 된다.

 

그리고, 위의 반례를 현명하게 다루는 것이 이번 문제의 핵심 아이디어이다.

 

1-2를 뽑는 것 보다 0-1과 2-3을 뽑는 것이 더 좋은 경우. 어떻게 처리해야 할까?

여기서 우리는 다음과 같은 천재같은 발상을 해낼 수 있다.

 

모든 간선들을 우선순위 큐에 넣어주는데, 항상 최소가 되는 것을 뽑아낸다. 다만 여기서 간선의 왼쪽과 오른쪽에 간선이 존재하면(위의 반례가 생길 여지가 있으면) 우리는 AB + CD - BC의 길이를 가진, 시작점과 끝점이 A와 D인 새로운 간선을 우선순위 큐에 넣어준다.

 

말로만 하면 이해가 안되니까 이를 다음의 게시판 반례로 보자.

 

input

6 3

1

3

4

6

7

9

 

output

6

 

처음에 우선순위 큐에는 5개의 간선이 모두 포함되어 있다.

간선은 (길이, 시작점, 끝점)의 튜플의 형태로 저장되어 있다.

이를 우선순위 큐에 넣으면 길이가 작은 순서대로 튀어나올 것이다.

 

처음으로 튀어나오는 간선은 1-2 간선이다.

위의 반례를 처리해주기 위해, 길이가 2+2-1 = 3이고 시작점과 끝점이 0, 3인 간선을 새로 추가해준다.

그리고 0-1 간선과 2-3 간선은 삭제해준다.

 

이를 K번 반복하면 된다(가장 맨 왼쪽이나 오른쪽 간선같이 반례가 안 발생하는 경우에는 그냥 그 간선을 빼주기만 한다).

 

그림으로 과정을 나타내면 아래와 같을 것이다.

최종 답은 6이다.

 

아이디어는 간단하다.

문제는 구현이다.

 

구현에서 가장 어려운 부분은 BC 간선을 선택했을 때, AB간선과 CD간선을 삭제하는 부분이다. 이를 우선순위 큐에서 ((O(N)))으로 remove했다가는 TLE가 발생할 것이 분명하다.

 

그래서 우리는 그래프 탐색에서 하듯이 초기값이 모두 False인 visited 배열을 만들어주고, 예를들어 1-2 간선을 선택했으면 visited[1]과 visited[2]는 True로 해준다.

그리고 최소 간선을 뽑았는데 start 또는 end 노드가 이미 방문한 지점(이미 연결됨)일 경우에는 continue를 이용해서 그냥 아무것도 안하고 넘어가준다. 이렇게 간선을 삭제하는 것이다!

 

그리고, 각 노드들의 바로 왼쪽/오른쪽 노드의 인덱스를 관리하는 리스트를 만들고, 왼쪽/오른쪽 노드까지 거리도 저장하는 리스트까지 해서 총 4개의 리스트를 만들어주면 우리는 헷갈리지 않고 쉽게 간선들을 관리해줄 수 있다.

 

예를 들어, 처음에 0-1-2-3이 있었는데 1-2가 선택당하고 0-3의 간선이 새로 큐에 추가된 경우, 이제 0의 오른쪽 노드는 3인 것이고 3의 왼쪽 노드는 0이 되는 것이다. 이를 관리해주는 리스트를 만드는 것이다.

 

이 아이디어들을 코드로 구현하면 아래와 같다.

 

import sys
import heapq

input = sys.stdin.readline
N, K = map(int, input().split())

arr = []

#--------4개의 리스트들 관리--------
left = [-1] + [i for i in range(N-1)]
lv = [0]
right = [i for i in range(1, N+1)]
rv = []
for i in range(N):
    arr.append(int(input()))

for i in range(1, N):
    lv.append(arr[i] - arr[i-1])

for i in range(N-1):
    rv.append(arr[i+1] - arr[i])
rv.append(0)
#-------------------------------

heap = []
for i in range(N-1):
    heapq.heappush(heap, (arr[i+1] - arr[i], i, i+1))

ans = 0
cnt = 0
visited = [False]*N
while cnt < K:
    if not heap:
        break
    #print(heap)
    l, start, end = heapq.heappop(heap)
    if visited[start] or visited[end]:
        continue

    #print(l, start, end)
    visited[start] = True
    visited[end] = True

    ans += l
    cnt += 1

    if left[start] >= 0 and right[end] < N:
        t = lv[start] + rv[end] - l
        heapq.heappush(heap, (t, left[start], right[end]))
        lv[right[end]] = t
        rv[left[start]] = t

        left[right[end]] = left[start]
        right[left[start]] = right[end]

print(ans)

 

Solved.ac 디스코드를 보니까 C++로 나이브하게 이 문제를 DP로 푼 사람도 있는 모양이다. 그분의 풀이에 따르면 이 문제는 골드 수준으로 다운그레이드 되어야 한다고......나중에 C++ 배우고 그렇게 한 번 풀어봐야겠다.