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

17124번 - 두 개의 배열

문제 이해

두 배열 A, B가 주어진다.
A의 각 요소에 대해 B의 값들 중 가장 차이가 적은(차의 절댓값) 값을 골라 새 배열 C를 만든다.
최종 C 배열의 값들의 합을 출력하면 되는 문제이다.

A배열의 길이는 N, B 배열의 길이는 M으로 주어진다.

접근

직접 배열을 만들 필요는 없으며, 합만을 변수로 관리하면 된다

• 브루트-포스
  A 배열의 N개 값에 대해 B 배열의 순회를 M번 진행해야 한다.
  이때 시간복잡도는 $O(NM)$ 이고, N과 M은 최대 $10^6$ 이므로, 1초에 약 1억번의 계산을 가정하면 시간초과를 받을것이 자명하다.

A의 각 원소에 대해 찾아야 하므로, N번 반복은 일어나야만 하는 반복이라 생각해 B에서 가장 가까운 값을 찾는 시간을 줄이는 방법으로 이분탐색을 생각했다.

• 이분탐색 - upper bound
  먼저 입력으로 주어진 B를 오름차순 정렬한다. - $O(MlogM)$
  A의 값을 하나 선택한 후, 그 값과 가장 가까운 B의 값을 선택할 때에, A에서 선택된 값보다 큰 값이 B에서 처음 등장하는 위치를 찾는다.
  즉, 해당 값에 대한 upper bound를 찾는다. - $O(logM)$
  찾은 upper bound l이라 한다면, 선택된 A[i]에 대해 B[l]과 B[l-1]중 차의 절댓값이 작은 값을 선택하면 된다.
  _{* l이 0이라면 B[0]을 선택}

이분탐색을 통한 구현의 시간복잡도는 $O((M+N)logM)$ 이다.

정답 코드

import sys
input = sys.stdin.readline
def getInts(): return map(int, input().split())


for _ in range(int(input())):
    N, M = getInts()
    A, B = [*getInts()], [*getInts()]
    B.sort()
    cnt = 0

    for a in A:
        l, r = 0, M-1
        while l < r:
            m = (l+r)//2
            if B[m] <= a:
                l = m + 1
            else:
                r = m
        if l == 0:
            cnt += B[0]
        else:
            cnt += B[l] if a - B[l-1] > B[l] - a else B[l-1]
    print(cnt)