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

숨·2024년 8월 13일

📌 문제 탐색하기

T : 테스트 케이스 수 $(1 ≤ T ≤ 10)$
n, m : 정수 배열 A의 길이, 정수 배열 B의 길이 $(1 ≤ n, m ≤ 10^6)$
A : n개의 정수 가지는 배열 $(1 ≤ 정수 ≤ 10^9)$
B : m개의 정수 가지는 배열 $(1 ≤ 정수 ≤ 10^9)$

A, B 배열을 이용해 C 배열을 만들고, C 배열 요소들의 총합을 구하는 문제이다.

문제 풀이

⭐️ 배열 C의 특징

len(C) = n

i번째 요소 = 배열 A의 i번째 요소와 가장 가까운 배열 B의 요소

위의 특징을 기반으로 C 배열을 만드는 코드를 아래와 같이 구현한다.
1️⃣ for문으로 배열 A의 요소를 순서대로 접근한다.
2️⃣ B의 요소 중 A의 해당 요소와 가장 가까운 요소를 이분 탐색을 통해 찾는다.
3️⃣ 찾은 요소를 C 배열 리스트에 추가한다.
4️⃣ C 배열의 총합은 sum() 함수를 이용해 구하고 출력한다.

가능한 시간복잡도

배열 A의 요소를 순서대로 접근 → $O(n)$
for문으로 A의 각 요소에 접근하면서 B에서 이분 탐색 → $O(n*log(m))$
C 배열 리스트에 요소 추가 → $O(n)$

최종 시간복잡도
총 $O(n*log(m))$ 이 된다.
최악의 경우 $O(10^6 log(10^6)) = O(6,000,000)$ 이 되는데
이는 2초 내에 연산이 가능하다.

알고리즘 선택

B의 요소 중 A의 해당 요소와 가장 가까운 요소를 이분 탐색으로 찾아내는 방법을 이용한다.

이분 탐색 구현

0️⃣ 테스트 케이스 별로 탐색하므로 함수로 따로 만들어준다.
1️⃣ 대소 비교를 통해 이분탐색 → 배열 B를 오름차순 정렬한다.
2️⃣ 시작값을 B의 첫번째 인덱스 0, 마지막값을 B의 마지막 인덱스 m-1로 지정한다.
3️⃣ 시작값이 마지막값과 같아지거나 더 작아질 때까지 탐색을 지속한다.
4️⃣ A[i]를 target, 정렬된 B의 중앙값 인덱스를 mid로 하여 대소 비교한다.
5️⃣ target = B[mid]라면? → 해당 값 바로 반환
6️⃣ target > B[mid]라면? → 시작값을 mid+1로 변경한다.
7️⃣ target < B[mid]라면? → 마지막값을 mid-1로 변경한다.
8️⃣ 얻은 start와 end 인덱스의 값 중 어떤 값이 절대값 차이가 가장 작은 값인지 계산해 반환

📌 코드 설계하기

이분 탐색 함수 구현
테스트 케이스 T 입력
n, m 입력
배열 A 입력
배열 B 입력
배열 C 리스트 초기화
for문으로 A의 각 요소 접근
6-1. 해당 요소와 가까운 B의 요소 찾기 위해 이분 탐색 수행
6-2. 찾은 요소를 C 리스트에 추가
C 리스트 합 출력

📌 시도 회차 수정 사항

1회차

def binary_search(target, B):
    # 해당 요소와 가까운 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 abs(B[start] - target) < abs(B[end] - target):
        return B[start]
    else:
        return B[end]

위와 같이 이분탐색을 구현했을 때 아래의 에러를 얻었다.
while문을 벗어나는 경우가 start <= end라는 조건이 지켜지지 않을 경우인데 이 부분을 제대로 고려해주지 않아서 발생한 문제이다.
start가 B의 크기를 넘어갔을 경우(end가 B의 가장 큰 값인데 start가 end보다 클 때), end가 0보다 작은 경우(start가 B의 가장 작은 값일 때)를 고려해주어야 한다.

📌 정답 코드

import sys

input = sys.stdin.readline


def binary_search(target, B):
    # 해당 요소와 가까운 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

    # start가 B의 크기보다 클 때
    if start >= m:
        return B[end]

    # end가 0보다 작을 때
    elif end < 0:
        return B[start]

    else:
        if abs(B[start] - target) < abs(B[end] - target):
            return B[start]
        else:
            return B[end]


# 1. 테스트 케이스 T 입력
T = int(input())

for _ in range(T):
    # 2. n, m 입력
    n, m = map(int, input().split())

    # 3. 배열 A 입력
    A = list(map(int, input().split()))

    # 4. 배열 B 입력
    B = list(map(int, input().split()))

    # 4-1. 배열 B 오름차순 정렬
    B.sort()

    # 5. 배열 C 리스트 초기화
    C = []

    # 6. for문으로 A의 각 요소 접근해 이분 탐색
    for i in range(n):
        closest = binary_search(A[i], B)
        # 6-1. 배열 C에 추가
        C.append(closest)

    # 7. C 리스트 합 출력
    print(sum(C))