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

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를 구하고 해당 배열의 모든 원소 합을 한 줄에 출력하시오.

코드

def bs(arr, target):
    s, e = 0, len(arr)-1
    res = -1
    while s <= e:
        mid = (s+e)//2
        if arr[mid] == target:
            res = mid
            break
        elif arr[mid] < target:
            res = mid
            s = mid+1
        else:
            e = mid-1
    return res


for _ in range(int(input())):
    n, m = map(int, input().split())
    a = list(map(int, input().split()))
    b = sorted(list(map(int, input().split())))
    sum_c = 0
    for i in range(n):
        near = bs(b, a[i])
        if near < m-1 and abs(b[near+1]-a[i]) < abs(b[near]-a[i]):
            near += 1
        sum_c += b[near]
    print(sum_c)

 

b배열에 있는 수 중에 a[i]보다 작으며 가장 가까운 값을 찾기 위해 이분탐색을 이용했다. 이분탐색을 이용하기 위해서는 b 배열이 오름차순으로 정렬되어 있어야 한다.

이분탐색을 수행하는 함수 bs()를 이용하여 a[i]보다 작은 수 중 가장 가까운 수를 찾아 near 변수에 저장한다. 그런데 우리는 가장 가까운 수를 찾는 것일뿐 그 수가 a[i]보다 작아야 한다는 조건은 없으므로 만약 b 배열의 near 다음 인덱스의 수가 a[i]에 더 가깝다면 near을 near+1로 갱신해준다. 이렇게 찾은 수는 c[i]이다. 이 수를 c 배열에 추가해주면 되는데 우리가 결국 구하는 건 c 배열의 총합이므로 c 배열을 따로 만드는 대신에 sum_c 변수를 만들어 여기에 하나씩 더해준다.