백준 1920 | 이진탐색 | 수 찾기 [파이썬 python]

1920번: 수 찾기

 

코드

# 이진 탐색 함수(재귀적 구현)
def binary_search(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2 # 중간점

    if array[mid] == target: # target을 찾았으면 리턴
        return mid
    elif array[mid] > target: # target이 중간점보다 작으면 왼쪽부분 탐색
        return binary_search(array, target, start, mid-1)
    elif array[mid] < target: # target이 중간점보다 크면 오른쪽부분 탐색
        return binary_search(array, target, mid+1, end)
    else:
        return None


n = int(input())
A = list(map(int, input().split()))
m = int(input())
B = list(map(int, input().split()))

A.sort() # A 정렬

# B의 요소가 A에 있는지 이진탐색으로 찾기
for i in B: 
    result = binary_search(A, i, 0, n-1)
    if result == None: # A에 없으면 0 출력
        print(0)
    else: # A에 있으면 1 출력
        print(1)

단순히 for문의 반복으로도 이 문제를 해결할 수는 있지만, 시간초과가 발생한다.

그래서 이진탐색을 사용하여 문제를 풀어야 시간 초과를 피할 수 있다.

 

먼저, 이진탐색 함수를 살펴보면 다음과 같다.

# 이진 탐색 함수(재귀적 구현)
def binary_search(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2 # 중간점

    if array[mid] == target: # target을 찾았으면 리턴
        return mid
    elif array[mid] > target: # target이 중간점보다 작으면 왼쪽부분 탐색
        return binary_search(array, target, start, mid-1)
    elif array[mid] < target: # target이 중간점보다 크면 오른쪽부분 탐색
        return binary_search(array, target, mid+1, end)
    else:
        return None

binary_search 매개변수로 (이진탐색을 수행할 정렬된 배열, 찾고자 하는 값, 시작점, 끝점)이 전달된다.

먼저, 시작점이 끝점보다 크면 이진탐색을 수행할 수 없으므로 None을 리턴하고, 그렇지 않은 경우 이진탐색을 수행한다.

 

이진탐색을 위해 중간점을 찾아야 한다.

중간점은 평균을 구하듯이 (시작점 + 끝점) // 2로 구한다. 나누기가 아닌 몫을 구하는 //을 사용한 이유는 mid는 인덱스로 사용되기 때문에 소수값은 될 수가 없기 때문이다.

 

중간점을 구했으면 본격적으로 이진탐색을 수행한다.

1. array[mid] == target이면 찾고자 하는 값을 찾은 것이므로 해당 인덱스 mid를 리턴한다.

2. array[mid] > target이면 target이 array[mid]보다 작은 값이므로, mid 기준 왼쪽 부분을 탐색해야 한다.
그러므로 binary_search함수를 다시 호출한다. 이때 매개변수로 전달되는 4가지 요소 중 end 값을 바꿔준다.
mid 기준 왼쪽 부분의 끝값의 인덱스는 mid-1이 된다.

3. array[mid] < target이면 target이 array[mid]보다 큰 값이므로, mid 기준 오른쪽 부분을 탐색해야 한다.
그러므로 binary_search함수를 다시 호출한다. 이때 매개변수로 전달되는 4가지 요소 중 start 값을 바꿔준다.
mid 기준 오른쪽 부분의 시작값의 인덱스는 mid+1이 된다.

4. 이진탐색을 통해 target을 찾았지만, 배열에 값이 존재하지 않으면 None을 리턴한다.

 

이제 문제에서 요구한 대로 출력을 하기 위한 코드를 작성한다.

# B의 요소가 A에 있는지 이진탐색으로 찾기
for i in B: 
    result = binary_search(A, i, 0, n-1)
    if result == None: # A에 없으면 0 출력
        print(0)
    else: # A에 있으면 1 출력
        print(1)

배열 B의 요소가 배열 A에 존재하는지 확인해야 하므로 B의 요소들을 for문으로 하나씩 불러온다.

해당 요소 i가 배열 A에 존재하는지 확인하기 위해 binary_search 함수를 수행한다.

처음 수행하는 함수의 시작점은 0번 인덱스가 되고, 끝점은 길이가 n인 배열의 마지막 인덱스가 되므로 n-1을 매개변수로 전달한다.

만약 리턴값이 None이면 배열 A에 i가 존재하지 않는 것이므로 0을 출력하고, 리턴값이 None이 아니면 배열 A에 존재하는 것이므로 1을 출력한다.