이진(이분) 탐색 공략법

문제

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

2343번: 기타 레슨

 

풀이

문제를 풀 때 항상 헷갈리는 문제 중 하나인 이진 탐색을 쉽게 푸는 방법에 대해 소개해보고자 한다.

위 문제를 예로 설명하겠음

 

문제 정리

문제를 요약하면 다음과 같다. 

1. 강의의 수 = N, 블루레이 개수 = M, 강의 길이 배열 = arr
2. M개의 블루레이에 강의를 나눠 담을 때 블루레이 크기의 최소값 구하기 = 정답

블루레이 크기를 평균값을 통해 유추하며 줄여나가 조건에 부합하는 최소값을 되기에

이진 탐색을 통해 문제를 풀 수 있다.

 

start, end 세팅

start = max(arr)
end = sum(arr)

이진 탐색을 위해 평균값을 구하기 위한 start, end 값을 세팅해야 하는데

• 최소한 하나의 강의가 들어갈 수 있는 블루레이 크기의 값을 start로 해야하니
  제일 긴 강의의 길이 max(arr)로 지정
• 1 <= M <= N의 조건이 명시되어 있으니 M = 1일 경우를 고려해
  모든 강의가 하나의 블루레이가 들어갈 수 있어서 end = sum(arr)로 지정

그 다음 while start <= end이 아닐 때까지 mid를 찾는다.

 

이진 탐색 기본틀

result = 0
while start <= end:
  mid = (start + end) // 2

이진 탐색 알고리즘은 항상 이렇게 시작한다고 보면 된다. 

제출할 결과물이 담길 result 값과, 

start가 end보다 커지면 이미 확인한 값들이기에 종료되게끔 while문 설정

mid 값은 start + end 값의 2로 나눈 나머지 뗀 값

 

풀이 핵심 1 - 문제 조건 반영

count = 1
temp = 0
for x in arr:
  if temp + x <= mid:
    temp += x
  else:
    temp = x
    count += 1

여기서 원하는 result 값을 찾아내기 위해 M과 비교하게될 count 값을 얻어야 한다.

코드작성에 앞서 말로 작성한다면 

1. 강의를 하나씩 꺼내 블루레이에 담다가
2. 강의 길이가 해당 블루에이 최대 저장 크기를 초과한다면 다음 블루레이에 저장

이 부분도 코드로 바로 작성하려고 하면 어느 시점에 count를 올리고 temp에 저장할지 헷갈리는 상황이 생긴다.

 

나는 이렇게 풀었다.

# arr = [3, 2, 1], mid = 5라고 했을 경우
for x in arr:  
  # x = 3일 때, temp = 3, count = 1
  # x = 2일 때, temp = 5, count = 1
  # x = 1일 때, temp = 1, count = 2

이걸 통해 미리 생각하고 갈 수 있는건 다음과 같다.

• count는 1부터 시작
• 새로 담을 강의를 넣었을 때 temp(현재 블루레이 저장량)가 안넘칠 경우 그대로 넣고 
  넘칠 경우 새 블루레이가 넣기 -> 넘칠 때만 count + 1

이렇게 차근차근 생각하면 코드를 바로 작성할 수 있게 된다.

 

풀이 핵심 2 - mid 값 조정

이제 얻어낸 count를 통해 mid 값을 조정하면 된다. (정확히는 start 또는 end 값 조정)

여기서도 조건을 하나씩 생각하면 되는데, 

1. count가 M과 같은 경우
2. count가 M보다 낮을 경우
3. count가 M보다 높을 경우

위 세가지 상황에 대해 start와 end 값을 조정하면서 최소값을 찾는다.

 

count가 M과 같은 경우

이 경우 result에 현재 mid 값을 담은 후

더 최소값을 구하기 위해 end = mid - 1을 세팅한다.

 

count가 M과 다른 경우

우리는 위에서 count를 얻는 과정을 통해 mid의 값 크기에 따라 count가 어떻게 변하는지 알 수 있었다.

• mid(블루레이 크기)가 크면 저장할 수 있는 길이도 길어져서 count는 줄어든다.
  mid up <-> count down
• 반대로 mid가 작으면 저장할 수 있는 길이가 줄어들어 count가 늘어난다.
  mid down <-> count up

 

위에서 count가 M과 같은 경우에 대한 end = mid - 1이 되야 하는 상황에 대해 고려한다면

1. end를 내리는 건 mid를 내리는 것
2. mid를 내리는 건 count를 올리는 것
3. count를 올리는 건 count가 M보다 낮다는 것

즉, count < M인 상황에선 end = mid - 1을 통해 내리는 결과로 이어지게 된다.

 

결과

if count <= M:
  result = mid
  end = mid - 1
else:
  start = mid + 1

따라서 위와 같이 작성하여 result, start, end 값을 조정한다.

 

정답

N, M = map(int, input().split())
arr = list(map(int, input().split()))

start = max(arr)
end = sum(arr)

result = 0
while start <= end:
  mid = (start + end) // 2
  count = 1
  temp = 0
  for x in arr:
    if temp + x <= mid:
      temp += x
    else:
      temp = x
      count += 1
  if count <= M:
    result = mid
    end = mid - 1
  else:
    start = mid + 1
    
print(result)

종합하면 위와 같다.

문제가 달라질 경우 보통 핵심 1 부분이 바뀌는데, 다른 알고리즘(ex. BFS 등)이 들어오기도 한다.

그에 대한 추가 문제로 아래 문제를 한번 풀어보는 것을 추천. 개인적으로 위 풀이 방법의 응용 문제로 너무 재밌게 푼 문제 중 하나이다.

 

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

1939번: 중량제한

 

이렇게 이진 탐색의 기본적인 풀이 방법에 대해 알아보았다.

 

 

* 위 내용은 알고리즘을 접한지 몇 달 되지 않은 초보가 스스로 정리한 글이기에 

이게 적용되지 않는 문제가 있거나 추천할 문제, 혹은 수정할 부분이 있으면 피드백! 부탁드립니다.