💻 Problem
정수 배열 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에 포함된 값들의 합을 구하는 프로그램을 작성하시오.
💡 Approach
A의 원소를 하나씩 살펴보며 B 배열에서 자신과 가장 가까운 값(절댓값 차이가 가장 작은 값)을 찾는다.
A[i]는 target이고 target과 가장 가까운 값을 B 배열에서 찾아야 한다.
A 배열과 B 배열의 최대 길이가 각각 10^6이므로 브루트포스로 풀면 시간 초과가 뜨게 된다.
이분탐색으로 구현해 보자.
A의 원소 하나하나씩 이분탐색으로 B에서 자신과 가장 가까운 값을 찾는다.
closest_idx, closest_val을 둬서 target과 가장 가까운 값의 인덱스와 값을 기록했다.
반복문을 활용해서 이분탐색을 구현했다.
이전보다 절댓값 차이가 더 작다면 closest_val 갱신하고, 만약 같다면 인덱스가 더 작을 때만 갱신한다.
타겟이 B[mid]보다 작으면 왼쪽을, 더 크면 오른쪽을 탐색한다.
기록해 둔 closest_idx가 이분탐색을 통해 찾은 인덱스이다.
B[closest_idx]를 리턴하면 된다.
✏️ Solution
import sys
input = sys.stdin.readline
def binary_search(target):
start, end = 0, m - 1
closest_idx, closest_val = m - 1, 10 ** 9 - 1
while start <= end:
mid = (start + end) // 2
# 이전보다 절대값 차이가 더 작다면 closest_val 갱신
# 같다면 더 작은 인덱스를 우선
compare_val = abs(target - B[mid])
if (closest_val > compare_val) or (closest_val == compare_val and closest_idx > mid):
closest_val = compare_val
closest_idx = mid
# 타겟이 더 작으면 왼쪽 탐색
if target < B[mid]:
end = mid - 1
# 타겟이 더 크면 오른쪽 탐색
else:
start = mid + 1
return B[closest_idx]
T = int(input())
for _ in range(T):
n, m = map(int, input().split())
A = list(map(int, input().split()))
B = sorted(map(int, input().split()))
res = 0
for a in A:
res += binary_search(a)
print(res)
⏱️ 시간복잡도
A 배열의 원소 개수: n (최대 10^6)
B 배열의 원소 개수: m (최대 10^6)
이분 탐색의 시간복잡도는 log N 이므로 이 코드의 시간복잡도는 n * log m이 된다.
10^6을 2의 제곱으로 치환하면 2^20이다.
대입하면 10^6 * 20 = 2 * 10^7이 된다.