[파이썬/python] 백준 - 17124 두 개의 배열

문제

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

---

문제 설명

1.   정수 배열 A와 B가 존재한다.
   1. A는 총 N개의 서로 다른 양의 정수를 포함한다.
   2. B는 총 M개의 서로 다른 양의 정수를 포함한다.
2. 배열 A와 B를 이용해서 배열 C을 만든다.
   1. C[i]는 배열 B에 있는 값중 A[i]에 가장 가까운 값으로 정의된다.
   2. 만약 해당 조건을 만족하는 값이 여러개 있는 경우, 그 중 가장 크기가 작은 값으로 정의된다.
3. 배열 C를 만들었을 때, 포함된 값들의 합을 구한다.

 

---

풀이

문제를 이해해보자.

A = [20,5,14,9], B = [16,8,12] 배열이 존재한다.

C[i]는 A[i]에 가장 가까운 B의 값이다. 20과 가장 가까운 값은 16이므로 C[1] = 16이 될 것이다.

 

나와 가장 가까운 값을 찾기 위해서는 어떻게 해야할까?

 

배열 B를 정렬한 후 처음부터 순차적으로 탐색하다가 나보다 큰 수가 나올 경우 해당 위치에서 내 앞과 뒤의 숫자를 확인하면 나와 가장 가까운 수를 찾을 수 있을 것이다.

 

하지만 배열 A와 B는 각각 $10^6$의 원소가 존재한다. 따라서 두 배열 모두 최대의 개수를 가질 때, A[i]를 찾기 위해서는 최악의 경우 $10^6$번의 연산이 필요하게 된다. 즉, 모든 C[i]를 찾기 위해서는 $N*M = 10^{12}$의 연산이 필요하다.

 

이 문제의 시간 제한은 1초($10^8$), 1억번의 연산을 넘게되면 시간초과가 발생한다. 즉 완전 탐색으로는 문제를 해결할 수 없다.

 

따라서 배열이 정렬되어 있을 때 내가 원하는 숫자를 찾는 이분탐색 알고리즘을 통해 문제를 풀이했다.

이분탐색을 사용하게 되면, A[i]와 가까운 수를 찾기 위해 걸리는 시간은 단 $log{10^6} = 20$에 불과하다.$(10^3 = 2^{10})$

 

따라서 약 $10^6 * 20 = 10*7$의 연산으로 문제를 풀이할 수 있다.

 

1.  배열 B를 오름차순으로 정렬한다.
2. A의 원소를 한 개씩 선택해서 B 배열에서 이분탐색을 진행한다.
   1. 반열린구간을 통해 이분탐색을 구현했기 때문에, e가 시작과 그대로라면 s위치의 값을 반환한다.
   2. 양쪽의 값의 차가 똑같다면, 절댓값이 더 작은(가까운) 값을 반환한다.
3. 합을 구하여 결과를 출력한다.

 

---

코드

import sys
input = sys.stdin.readline

def binary_search(target):
    s = 0
    e = len(B)
    while s + 1 < e:
        mid = (s + e) // 2
        if B[mid] <= target:
            s = mid
        else:
            e = mid
    if e == len(B):
        return B[s]
    else:
        if abs(B[s]-target) <= abs(B[e]-target):
            return B[s]
        else:
            return B[e]

for tc in range(int(input())):
    n,m = map(int,input().split())
    A = list(map(int,input().split()))
    B = sorted(list(map(int,input().split())))
    res = 0
    for num in A:
        res += binary_search(num)
    print(res)

---

시간복잡도

• 시간복잡도는 $O(N*logM)$이다. 따라서 약 $O(10^7)$의 시간복잡도가 소요된다.