구현

이 장에서는 탐색의 기본 부품인 lower_bound를 Python으로 깔끔하게 구현하고 한 줄씩 설명합니다. 이어서 같은 함수를 C++, Java, TypeScript로 옮기고, 언어마다 이미 들어 있는 표준 함수와 가운데 인덱스를 계산할 때의 오버플로 문제를 짚습니다.

Python 구현

def lower_bound(a, x, lo=0, hi=None):
    if hi is None:
        hi = len(a)
    while lo < hi:
        mid = (lo + hi) // 2
        if a[mid] < x:
            lo = mid + 1
        else:
            hi = mid
    return lo


def upper_bound(a, x, lo=0, hi=None):
    if hi is None:
        hi = len(a)
    while lo < hi:
        mid = (lo + hi) // 2
        if a[mid] <= x:
            lo = mid + 1
        else:
            hi = mid
    return lo


def contains(a, x):
    i = lower_bound(a, x)
    return i < len(a) and a[i] == x

한 줄씩 보면 다음과 같습니다.

• lo=0, hi=None: 탐색할 반열린 구간 [lo, hi)입니다. 기본값은 배열 전체이며, 일부 구간만 볼 수도 있게 bisect 모듈과 같은 매개 변수를 둡니다.
• while lo < hi: 구간에 후보가 하나라도 남아 있는 동안 반복합니다. lo == hi가 되면 그 자리가 답입니다.
• mid = (lo + hi) // 2: 가운데 인덱스입니다. Python 정수는 크기 제한이 없어서 더해도 넘치지 않습니다.
• if a[mid] < x: lo = mid + 1: mid의 값이 x보다 작으면 mid와 그 왼쪽은 답이 될 수 없으므로 버립니다.
• else: hi = mid: mid의 값이 x 이상이면 mid 자체가 답일 수 있으므로 mid는 남기고 오른쪽만 버립니다.
• return lo: 끝났을 때 lo는 a[i] >= x인 첫 인덱스이며, 그런 원소가 없으면 len(a)입니다.

upper_bound는 비교 하나(<를 <=로)만 다릅니다. contains는 lower_bound 위치가 배열 안에 있고 그 값이 정확히 x인지 확인합니다. 인덱스 범위를 먼저 검사해야 IndexError가 나지 않습니다.

표준 라이브러리와 맞춰 보기

직접 구현했다면 표준 구현과 무작위로 비교해 보는 것이 가장 확실한 검증입니다.

import random
from bisect import bisect_left, bisect_right

for _ in range(10_000):
    a = sorted(random.randint(0, 20) for _ in range(random.randint(0, 15)))
    x = random.randint(-1, 21)
    assert lower_bound(a, x) == bisect_left(a, x)
    assert upper_bound(a, x) == bisect_right(a, x)
    assert contains(a, x) == (x in a)
print("ok")

빈 배열, 모든 원소보다 작거나 큰 값, 전부 같은 값처럼 경계 사례가 무작위 입력에 자연스럽게 섞여 들어갑니다. Python 3.10부터는 bisect_left(records, 30, key=lambda r: r.age)처럼 key도 줄 수 있습니다.

C++

#include <cstddef>
#include <vector>

std::size_t lower_bound_index(const std::vector<int>& a, int x) {
    std::size_t lo = 0, hi = a.size();
    while (lo < hi) {
        std::size_t mid = lo + (hi - lo) / 2;
        if (a[mid] < x) lo = mid + 1;
        else hi = mid;
    }
    return lo;
}

C++ 표준 라이브러리에는 std::lower_bound(a.begin(), a.end(), x)가 있어 반복자를 돌려줍니다. 인덱스가 필요하면 a.begin()을 빼면 됩니다. std::upper_bound, 두 경계를 함께 주는 std::equal_range, 조건의 경계를 찾는 std::partition_point도 있습니다.

Java

public final class Search {
    public static int lowerBound(int[] a, int x) {
        int lo = 0, hi = a.length;
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (a[mid] < x) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

Java의 Arrays.binarySearch는 값을 찾으면 그 인덱스를, 못 찾으면 -(끼워 넣을 자리) - 1을 돌려줍니다. 중복이 있으면 어느 위치를 돌려줄지 보장하지 않으므로, 첫 위치가 필요하면 위와 같이 lower_bound를 직접 쓰는 편이 안전합니다.

TypeScript

export function lowerBound(a: readonly number[], x: number): number {
  let lo = 0;
  let hi = a.length;
  while (lo < hi) {
    const mid = Math.floor((lo + hi) / 2);
    if (a[mid] < x) lo = mid + 1;
    else hi = mid;
  }
  return lo;
}

console.log(lowerBound([1, 2, 2, 2, 3, 5], 2)); // 1

JavaScript에는 배열용 표준 이진 탐색 함수가 없어서 이렇게 직접 작성하는 일이 많습니다. 숫자가 64비트 부동소수점이라 lo + hi가 넘칠 걱정은 없지만, (lo + hi) >> 1은 32비트 부호 있는 정수로 바뀌므로 아주 큰 값에서는 틀릴 수 있습니다. 문자열 배열이라면 비교 함수를 매개 변수로 받도록 일반화합니다.

가운데 인덱스와 오버플로

C, C++, Java처럼 정수 크기가 고정된 언어에서 (lo + hi) / 2는 lo와 hi가 모두 클 때 더하는 순간 넘칠 수 있습니다. Java의 Arrays.binarySearch에도 실제로 이 버그가 오래 남아 있다가 2006년에 고쳐졌습니다. 안전한 방법은 두 가지입니다.

• lo + (hi - lo) / 2: 차이를 반으로 나눠 더하므로 hi를 넘지 않습니다.
• Java의 (lo + hi) >>> 1: 부호 없는 오른쪽 시프트라 합이 음수로 넘쳐도 올바른 값이 나옵니다(lo와 hi가 0 이상일 때).

정리

• 반열린 구간, lo = mid + 1과 hi = mid 두 가지 갱신만 쓰는 lower_bound를 기본 부품으로 삼습니다.
• upper_bound는 비교 하나만 바꾸며, 값 존재 여부는 lower_bound 결과로 확인합니다.
• 언어마다 bisect, std::lower_bound, Arrays.binarySearch 같은 표준 구현이 있으니 반환값 규칙을 확인하고 씁니다.
• 고정 크기 정수 언어에서는 lo + (hi - lo) / 2로 오버플로를 막습니다.