np.searchsorted는 이진 탐색 알고리즘을 구현한 NumPy 함수다. 히스토그램 계산과 구간 조회 등에 쓰이며, NumPy 2.5에서는 벤치마크 기준으로 최대 25배 빨라졌다. 이 최적화는 SciPy와 scikit-learn 등 searchsorted에 의존하는 라이브러리에도 영향을 줄 수 있다.

이진 탐색 변형과 배치 처리의 효과와 Algorithmica의 이진 탐색 사례 연구는 분기 제거, 배치 처리, 캐시 친화적 데이터 배치 같은 기법을 살펴본다. 이 글은 이런 아이디어를 NumPy의 벡터화 연산으로 표현한 뒤, 알고리즘을 NumPy에 적용한 과정을 설명한다.

searchsorted는 Python Array API 표준에도 포함돼 있어, 배열 라이브러리마다 같은 연산을 어떻게 구현하고 병렬성을 활용하는지 비교할 수 있다.

벤치마크 조건과 벡터화

문제는 정렬된 고정 배열과 검색 키 목록이 주어졌을 때, 각 키를 삽입할 위치를 찾는 것이다. 기존 방식은 키마다 이진 탐색을 순차적으로 수행한다. 글에서 비교에 사용한 NumPy 2.4 구현도 기본적으로 키마다 독립적인 이진 탐색을 한다.

벤치마크에는 균등 분포의 무작위 np.int32 배열 두 개를 사용했다. 검색 키는 10,000개로 고정하고 정렬 배열의 길이는 최대 $2^{30}$($4\ GiB$)까지 늘렸다. 두 배열은 메모리에 연속 배치했다. 각 벤치마크를 50회 반복해 최솟값을 보고했다. 실험 환경은 메모리 32 GB를 갖춘 Apple M1 Pro 탑재 MacBook Pro다. M1 Pro의 성능 코어당 L1 데이터 캐시는 128 KB, 성능 코어가 공유하는 L2 캐시는 12 MB다.

여러 검색을 한 번에 진행하도록 각 검색의 lo와 hi 상태를 배열로 저장하고, NumPy 연산으로 동시에 갱신할 수 있다. 배열 연산은 컴파일된 C++ 루프에서 실행되므로 Python 반복문의 오버헤드를 줄인다. 다만 검색마다 수렴 시점이 다를 수 있어, 첫 벡터화 구현은 아직 진행 중인 검색을 나타내는 active 마스크를 관리한다.

다음 단계에서는 검색 구간이 수렴한 뒤에도 이후 반복에서 그대로 유지되도록 갱신 방식을 바꾼다. 그러면 모든 검색이 np.ceil(np.log2(n))회 안에 끝나고 active 추적을 없앨 수 있다. 글에서는 이 변경만으로 최대 2배 빨라졌다고 설명한다. 이 간단한 구현은 배열 a가 비어 있지 않다고 가정한다. 비슷한 형태는 JAX의 scan 기반 구현에서도 볼 수 있다.

NumPy 2.5에 적용한 최적화

벡터화 구현은 검색마다 두 배열을 유지하므로, 이를 그대로 네이티브 코드로 옮기면 검색 키 수에 비례하는 추가 메모리가 필요하다. 대신 각 구간을 왼쪽 경계와 길이로 표현하고, 모든 검색 구간의 길이가 반복마다 같도록 구성하면 오른쪽 경계를 별도로 저장하지 않아도 된다. 결과 배열인 base만 검색별 상태로 사용하고, 구간 길이는 하나의 공통 값으로 관리한다. 이 방식은 결과 배열 외에 추가로 필요한 메모리가 O(1)이다.

이 재구성은 Algorithmica 사례 연구의 분기 없는 이진 탐색과 관련이 있다. 구간을 [base, base + length]로 표현하므로 length = 1일 때 삽입 위치를 확정하는 마지막 단계가 필요하다.

이 알고리즘은 NumPy의 C++ 구현으로 옮겨져 PR #30517을 통해 NumPy 2.5 릴리스에 포함됐다. 첫 반복에서는 모든 키가 같은 중앙값과 비교되므로 중앙값을 한 번만 읽는다. 결과 항목의 초기값이 암묵적으로 0인 점도 이용해 첫 반복에서 초기화 쓰기와 읽기를 생략했다.

벤치마크에서 NumPy 2.5의 네이티브 구현은 NumPy 2.4보다 최대 25배 빨랐다. 벡터화한 Python 구현과 비교하면 작은 배열에서 최대 2배 빠를 수 있으며, 배열이 커질수록 차이는 줄었다. 메모리 측면에서는 C++ 구현이 검색 키 수에 비례하는 상태 배열을 추가로 만들지 않고 O(1) 메모리만 사용한다. 벡터화 Python 구현은 검색 상태를 저장할 배열이 더 필요하다.

다른 배열 라이브러리와의 비교

비교 대상은 JAX, TensorFlow, PyTorch다. JAX와 NumPy는 벡터화·배치 연산으로 메모리 지연을 감추는 방식이고, TensorFlow와 PyTorch는 검색 키를 나눠 CPU 스레드에서 병렬 처리한다. 구현 세부사항은 PyTorch와 TensorFlow 코드에서 확인할 수 있다.

비교 실험에서는 병렬 실행을 8개 코어로 제한하고 검색 키를 10,000개에서 20,000개로 늘려, 다중 스레드 구현이 스레드 스케줄링 비용을 상쇄할 만큼 충분한 작업을 수행하도록 했다. 이 벤치마크에서 NumPy는 선택된 라이브러리들과 경쟁력 있는 성능을 보였으며, 검색 배열이 CPU 캐시를 넘어선 뒤에는 모든 구현이 비슷한 양상을 보였다. 다중 스레딩을 끄면 PyTorch와 TensorFlow의 성능이 떨어지고 NumPy 2.4와 비슷한 추세를 보였다.

글은 스레드마다 검색을 배치하는 기법을 결합할 수 있는지 추가로 측정해 볼 만하다고 제안한다. 다만 메모리 하위 시스템이 포화되면 코어를 늘려도 메모리 대역폭을 놓고 경쟁할 수 있다. 대안으로 Eytzinger 레이아웃 같은 데이터 배치 방식도 언급한다. 이 레이아웃은 정렬된 표현이 아니므로, 이를 searchsorted 인터페이스로 제공하려면 API 의미를 신중하게 정의해야 한다는 설명이다.

벤치마크 결과는 해당 실험 조건에서 나온 수치다. NumPy의 최적화는 SciPy와 scikit-learn 등 의존 라이브러리에 도움이 될 수 있으며, 자체 이진 탐색을 구현한 다른 라이브러리도 배치 처리 기법을 적용해 볼 수 있다.