백준20551 - Sort 마스터 배지훈의 후계자 (이분탐색) Python

https://www.acmicpc.net/problem/20551

해당 문제는 주어진 배열의 길이 | 찾아야 할 요소 길이 | 배열 요소 | 찾아야 할 요소를 주고, 

해당 배열을 정렬한 뒤, 찾아야 할 요소가 정렬된 배열에 존재하면 해당 요소의 index를, 없으면 -1 을 반환하는 문제다.

 

이분탐색 문제인데 2분만에 풀진 못했다.

 

그 이유로는 처음 완성한 이진탐색 로직에는, 배열에 중복된 요소가 없으리라 가정하고 작성되었기 때문에

만약 해당 함수에 arr = [5,5,5,5,5], x = 5 가 인자로 들어오더라도 

가장 처음의 index값이 아닌, 중간값 index를 반환한다는 문제를 발견했다.

def bisect (arr, x) :
  left, right = 0, len(arr)
  if arr[left] > x or x > arr[right-1] : 
    return -1
  
  while left <= right :
    mid = (left + right) // 2
    if arr[mid] == x :
      return mid
    elif arr[mid] > x :
      right = mid - 1 
    else : 
      left = mid + 1
      
  return -1

 

이 문제를 찾기 위해 구글을 찾다가 이진탐색에서 lower_bound 라는 로직 개념도 알게 되었는데, 

이를 활용하여 수정한 로직이 아래 로직과 같다.

def bisect (arr, x) :
  if arr[0] > x or x > arr[-1] : 
    return -1
  
  left, right = 0, len(arr) - 1
  
  while left <= right :
    mid = (left + right) // 2
    if arr[mid] > x :
      right = mid - 1 
    elif x > arr[mid] : 
      left = mid + 1
    elif arr[mid] == x :
      if right == mid : break
      else : right = mid
  
  return mid if arr[mid] == x else -1

 

즉, if arr[mid] == x 인 값을 찾더라도 right와 mid가 같지 않으면 해당 로직을 다시 반복하여(좌측으로 범위를 줄이며)

조건을 만족하는 가장 작은 index를 찾는 로직으로 바꾸었다.

 

전체 소스 코드

import sys
data_temp = sys.stdin if sys.platform == 'linux' else open('입력.txt', 'r')
input_data = data_temp.read().splitlines()

def solution (data) :
  n,_ = map(int, data.pop(0).split())
  arr = sorted(list(map(int, data[:n])))
  catch = list(map(int, data[n:]))
  result = [bisect(arr, v) for v in catch]
  print('\n'.join(map(str, result)))

def bisect (arr, x) :
  if arr[0] > x or x > arr[-1] : 
    return -1
  
  left, right = 0, len(arr) - 1
  
  while left <= right :
    mid = (left + right) // 2
    if arr[mid] > x :
      right = mid - 1 
    elif x > arr[mid] : 
      left = mid + 1
    elif arr[mid] == x :
      if right == mid : break
      else : right = mid
  
  return mid if arr[mid] == x else -1    
  
solution(input_data)