[Python] 백준 17124. 두 개의 배열 +) 이분탐색 정리!

문제

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

 

정수 배열 A 와 B가 있다. A는 총 n개의 서로 다른 양의 정수를 포함하고 B는 총 m개의 서로 다른 양의 정수를 포함한다.

A, B를 이용해서 길이가 n인 새로운 배열 C를 만들어보자.

• C[i] 는 배열 B에 있는 값 중 A[i] 에 가장 가까운 값 (절대값 차이가 가장 작은 값)으로 정의 된다. 
• 만약 이 조건을 만족하는 값들이 여럿 있는 경우, 그 중 가장 크기가 작은 값으로 정의 된다.

예를 들어 A = [20, 5, 14, 9] 그리고 B = [16, 8, 12] 라고 해보자.

• C[1] = 16 이다 - 왜냐하면 B[1] = 16이 A[1] = 20에 가장 가깝기 때문이다.
• C[2] = 8 이다 - 왜냐하면 B[2] = 8이 A[2] = 5에 가장 가깝기 때문이다.
• C[3] = 12 이다 - 왜냐하면 B[1] = 16 와 B[3] = 12 모두 A[3] = 14에 가장 가깝지만, B[3]의 값이 더 작기 때문이다.
• C[4] = 8이다.

이 예제의 경우 C = [16, 8, 12, 8]으로 정의된다.

두 배열 A와 B가 주어졌을 때, 새로운 배열 C를 계산하여 배열 C에 포함된 값들의 합을 구하는 프로그램을 작성하시오.

입력

첫 줄에 테스트 케이스의 수 T (1 <= T <= 10)가 주어진다.

각 테스트 케이스는 세 줄에 걸쳐서 주어진다.

첫 줄에는 n과 m이 공백으로 구분되어 주어진다 (1 <= n, m <= 10^6).

두 번째 줄에는 공백으로 구분된 n개의 정수가 주어지며, A[1] 부터 A[n]을 나타낸다 (각각의 값은 1이상 10^9 이하이다).

세 번째 줄에는 공백으로 구분된 m개의 정수가 주어지며, B[1] 부터 B[m]을 나타낸다 (각각의 값은 1이상 10^9 이하이다).

앞서 언급한대로, A와 B는 각각 서로 다른 양의 정수들을 포함한 배열들이다.

출력

각 테스트 케이스에 대해 배열 C를 구하고 해당 배열의 모든 원소 합을 한 줄에 출력하시오.

---

풀이

이 문제는 딱 보자마자 이분탐색을 생각했어야 했는데... 못했다. 일단 이분탐색을 잘 못해서 그런 것 같아 이번에 겸사겸사 이분탐색을 확실히 정리를 하고 문제 해설을 할까 한다. 일단 문제 해설부터~

 

문제는 A배열에 있는 값들과 B배열에 있는 값들을 비교해서 가장 가까이에 있는 값을 대신 넣어주고 합을 출력하는 내용이었다. 

< 문제 이해하기 쉬운 예제 >

예를 들어 A = [20, 5, 14, 9] 그리고 B = [16, 8, 12] 라고 해보자.

• C[1] = 16 이다 - 왜냐하면 B[1] = 16이 A[1] = 20에 가장 가깝기 때문이다.
• C[2] = 8 이다 - 왜냐하면 B[2] = 8이 A[2] = 5에 가장 가깝기 때문이다.
• C[3] = 12 이다 - 왜냐하면 B[1] = 16 와 B[3] = 12 모두 A[3] = 14에 가장 가깝지만, B[3]의 값이 더 작기 때문이다.
• C[4] = 8이다.

이 예제의 경우 C = [16, 8, 12, 8]으로 정의된다.

그러니까 A[i] 와 B 전체를 한번씩 탐색해서 가장 가까운 수를 C[i]에 넣어주고 이를 합해주면 되는 것이었다. 이때 B 전체를 탐색하면서 A[i]와 비교하는 걸 이분탐색으로 해결하면 되는 것이다!

import sys
input = sys.stdin.readline

def binary_search(target, b):
    start, end = 0, m - 1

    while start <= end:
        mid = (start + end) // 2
        if target == b[mid]:
            return b[mid]
        elif target > b[mid]:
            start = mid + 1
        else:
            end = mid - 1

    if start >= m:
        return b[end]

    elif end < 0:
        return b[start]

    else:
    	diff_start = abs(b[start] - target)
    	diff_end = abs(b[end] - target)
    
    	if diff_start < diff_end:
        	return b[start]
    	elif diff_start > diff_end:
        	return b[end]
    	else:
        	return min(b[start], b[end])
        
t = int(input().strip())

for _ in range(t):
    n, m = map(int, input().rsplit())
    a = list(map(int, input().rsplit()))
    b = list(map(int, input().rsplit()))
    b.sort()
    c = []
    
    for i in range(n):
        close = binary_search(a[i], b)
        c.append(close)

    print(sum(c))

 

이분탐색 부분은 다음 블로그를 참고했다. 

https://velog.io/@klaus/%ED%8C%8C%EC%9D%B4%EC%8D%ACPython-%EB%B0%B1%EC%A4%80-17124%EB%B2%88-%EB%91%90-%EA%B0%9C%EC%9D%98-%EB%B0%B0%EC%97%B4

[파이썬/Python] 백준 17124번: 두 개의 배열

사실상 이분탐색만 잘 사용하면 나머지는 입력받고 출력하는 정도이기 때문에... 코드 전체에 도움을 받았다고 해도 과언은 아닌 듯하다. 

 

이분탐색 파트를 뜯어보면서 이해를 해보자면...

일단 A[i]에 대해 B 배열을 전부 확인해봐야하는데, 진짜 다 확인해보면 시간이 너무 오래걸리므로

   b = list(map(int, input().rsplit()))
    b.sort()
    c = []
    
    for i in range(n):
        close = binary_search(a[i], b)

1. b배열을 입력받고 정렬을 해준다. (크기 순으로 절반 나눠서 탐색할 거기 때문에)

2. 그러고 for문으로 A배열의 값을 하나씩 골라주고

3. 이분탐색 함수를 돌려준다. 

def binary_search(target, b):
    start, end = 0, m - 1

    while start <= end:
        mid = (start + end) // 2
        if target == b[mid]:
            return b[mid]
        elif target > b[mid]:
            start = mid + 1
        else:
            end = mid - 1

    if start >= m:
        return b[end]

    elif end < 0:
        return b[start]

    else:
    	diff_start = abs(b[start] - target)
    	diff_end = abs(b[end] - target)
    
    	if diff_start < diff_end:
        	return b[start]
    	elif diff_start > diff_end:
        	return b[end]
    	else:
        	return min(b[start], b[end])

4. 이분탐색 함수에서는 일단 시작과 끝을 더해주고 2로 나눠줘서 중간을 찾는다

5. B의 중간 숫자가 A[i](즉, target)이랑 같으면 => return b[mid]

target이 더 크면 중간숫자보다 큰 숫자들 중 있는 것이므로 start 를 mid + 1해주고

target이 더 작으면 중간숫자보다 작은 숫자들 중에 가장 가까운 숫자가 있는 것이므로 end를 mid-1해준다.

6. 그리고 이렇게 바뀐 start, end를 end가 start보다 작아지기 전까지 계속 반복...

 

7. 그렇게 구한 start가 m보다 크거나 같으면 즉, start가 배열의 범위를 넘어갔다면 b[end]로 출력

end가 0보다 작으면, end가 배열의 범위를 넘어갔다는 뜻이므로 b[start]를 출력

 

8. 만약 어디에도 포함하지 않고 가장 가까운 것이 두 가지라면, 더 작은 것을 출력해야하므로 두개를 비교 후 작은 것을 출력

 

시간복잡도

O(T(n + m) log m) 일 것 같다! 이분탐색이 log m이라고 한다..