[BOJ/Python] 2230_수 고르기

문제 출처:https://www.acmicpc.net/problem/2230

2230번: 수 고르기

문제

N개의 정수로 이루어진 수열 A[1], A[2], …, A[N]이 있다. 이 수열에서 두 수를 골랐을 때(같은 수일 수도 있다), 그 차이가 M 이상이면서 제일 작은 경우를 구하는 프로그램을 작성하시오.

예를 들어 수열이 {1, 2, 3, 4, 5}라고 하자. 만약 M = 3일 경우, 1 4, 1 5, 2 5를 골랐을 때 그 차이가 M 이상이 된다. 이 중에서 차이가 가장 작은 경우는 1 4나 2 5를 골랐을 때의 3이 된다.

입력

첫째 줄에 두 정수 N, M이 주어진다. 다음 N개의 줄에는 차례로 A[1], A[2], …, A[N]이 주어진다.

출력

첫째 줄에 M 이상이면서 가장 작은 차이를 출력한다. 항상 차이가 M이상인 두 수를 고를 수 있다.

제한

• 1 ≤ N ≤ 100,000
• 0 ≤ M ≤ 2,000,000,000
• 0 ≤ |A[i]| ≤ 1,000,000,000

예제 입력 1

3 3
1
5
3

예제 출력 1

4

---

풀이

 

두 수를 골랐을 때, 두 수의 차가 M이상이면서 가장 작은 값을 찾아내는 문제이다.

우선 이 문제를 봤을 때 생각나는 방법이 두 가지가 있었다.

1. itertools 모듈에 있는 combinations(조합)을 이용한다.
2. 주어진 정수들 중 두 개를 선택하는 모든 경우를 찾아낸다. (nC2)
3. 모든 경우의 차이를 구하여 M이상이면서 가장 작은 값을 구한다.

 

1. 우선 리스트를 오름차순으로 정렬한다.
2. 두개의 포인터를 이용하여 탐색한다. (왼쪽 포인터와 오른쪽 포인터)
3. 왼쪽 포인터는 고정시키고, 오른쪽 포인터를 하나씩 증가시키며 두 포인터가 가리키는 값의 차이를 구한다.
4. 차이가 M보다 작다면 오른쪽 포인터를 하나씩 증가시킨다.
5. 차이가 M이상이라면, min(지금까지 M이상이면서 가장 작은 차이, 지금 구한 차이)를 통하여 M이상이면서 가장 작은 값을 구한다.
6. 왼쪽 포인터를 하나 증가시키고, 오른쪽 포인터는 왼쪽 포인터의 위치부터 3~5의 과정을 반복한다.

 

우선 조합을 이용하는 방법의 경우, 모든 경우를 구해야 하기 때문에 시간이 오래 걸리게 된다. 따라서 투 포인터를 이용하는 두 번째 방법을 이용하여 코드를 작성하였다.

하지만, 결과는 시간초과가 나왔고, 그 이유를 생각해보니 다음과 같은 결론을 내릴 수 있었다.

만약 리스트에 1~100까지의 정수가 존재하고, M=99라고 하자.

이 경우, 첫번째 1과 다른 수의 차이가 M이상인 경우를 구하기 위해서, 오른쪽 포인터는 1부터 100까지 100번을 이동해야 한다. 2의 경우 2부터 100까지 99번을 이동해야 하고, 3의 경우 98번을 이동해야 한다.

따라서 이는 O(N**2)의 시간복잡도를 가지게 된다.

사실 이 경우, 투 포인터를 구현한 것 같지만, 이중 for문을 구현한 것과 같다.

 

다시 1의 경우를 생각해 보면 결국 M=99이므로 100 이상의 수를 찾으면 되는 것이고, binary search를 이용하여 어떤 수 num에 대하여 차이가 M이상이 나는 수,  즉 num+M 이상의 수를 찾는 코드를 작성하여 문제를 해결하였다. 

코드는 다음과 같다.

 

 

import sys
import bisect
 
n, m = map(int,sys.stdin.readline().split())
n_list = []
for _ in range(n):
    n_list.append(int(sys.stdin.readline()))
 
n_list.sort()
diff = 2*n_list[-1]
 
for start in range(n):
 
    idx = bisect.bisect_left(n_list,m+n_list[start])    
    if idx==n:
        continue
 
    diff = min(diff,n_list[idx] - n_list[start])
    if diff==m:
        break
 
 
print(diff)

• bisect_left()의 경우, m+n_list[start]가 n_list에 몇 번째 index의 들어가는 지를 알려준다. (같은 수가 있을 경우 제일 왼쪽 index)
• 반대로 생각하면 m+n_list[start] 이상의 수 중 가장 작은 수의 index를 구할 수 있다. (즉, 두 수의 차가 m 이상이 되게 하는 가장 작은 수)
• 여기서 두 수의 차가 m 이상이 되게 하는 수가 없다면 idx의 결과는 n이 나온다. (즉 제일 오른쪽 값)
• 이 경우, continue를 이용하고, 아닌 경우 두 수의 차를 구하여 가장 작은 지 확인한다.
• 만약 차이가 m이라면, 탐색을 종료한다. (m 이상인 수 중 가장 작은 값을 구해야 하는데 m이므로 더 이상 탐색할 이유가 없다.)

bisect의 경우 아래와 같이 생각하면 된다. binary search 이므로 오름차순으로 정렬되어있어야 한다.

오름차순으로 정렬된 n_list에서 5가 들어갈 수 있는 위치는 3번째 index이다.

 

만약 이처럼 같은 수가 있을 때, bisect_left의 경우 가장 왼쪽 index이고, bisect_right의 경우 가장 오른쪽 index이다.

다시 말해 8이 리스트에 들어간다면, 가장 왼쪽의 들어갈 경우 bisect_left의 결괏값이 그 index인 것이고, 가장 오른쪽의 들어갈 경우 bisect_right의 결괏값이 그 index인 셈이다.

 

이 경우 12가 들어가는 위치는 7번 index가 된다.