배열 내에 k번째로 작은 수 구하기

어떤 배열이 주어졌을 때, 그 배열 속에서 k번째로 작은 수를 구하는 일은 심심치 않게 발생한다. 나 같은 경우에는, 어떤 확률을 나타내고 있는 n차원 백터내에서 k번째로 큰 확률을 가진 클래스를 찾고 싶었다. 그럼 이를 어떻게 구현할 수 있을까? 가장 간단하게 떠오르는 방법은 (1) 주어진 벡터를 오름차순으로 정렬한 뒤 (인덱스를 알고싶다면 인덱스도 함께 저장), (2) k번째 값을 찾는 것이다. 해당 방법은 O(nlogn)의 시간복잡도를 가지며 적당한 시간 내에 끝나기 때문에 실용적이다. 하지만, 이것보다 더 빠르게 구현할 수 있는 방법이 있을까?

검색 알고리즘 (selection algorithm)

이번에 찾아보면서 알게 되었는데 배열 속에서 k번째로 작은 수를 구하는 문제는 알고리즘 분야에서는 꽤나 유명하고 이미 충분히 연구된 문제인 것 같다. 해당 문제를 푸는 알고리즘으로 앞서 말한 정렬을 하고 풀거나, heap 자료구조를 활용해서 푸는 방법이 간단하며 둘다 O(nlogn)의 시간복잡도를 가진다. 하지만, 정렬 중에서 퀵정렬 (quicksort)의 내부 구조를 조금더 활용하면 평균시간복잡도가 O(n)으로 줄어든 알고리즘을 고안할 수 있다.

퀵선택 (quickselect)

퀵선택 알고리즘은 퀵정렬 알고리즘을 정렬이 아닌 검색으로 사용하는 알고리즘이다. 퀵정렬 알고리즘은 먼저, 축(pivot)이 되는 값을 설정하고, 그 축에 해당하는 값을 기준으로 그 값보다 작은 값들로 이루어진 집합 (A)과 그 값보다 큰 값들로 이루어진 집합 (B)으로 배열을 나눈다. 두 배열로 나눈 상태에서, 정렬이 아닌 배열 내에서 k번째로 작은 수를 찾으려면 어떻게 해야할까?

 

만약에 축에 해당하는 값보다 작은 값으로 이루어진 집합(A)의 크기가 k보다 크거나 같으면, k번째로 작은 수는 A안에 있을 것이기 때문에 문제는 A내에서 k번째로 작은 수를 푸는 문제로 바뀌게 된다. 만약 A의 크기가 k보다 작다면, 찾고자 하는 값은 B안에 있을 것이기 때문에, 문제는 B내에서 k-C(A)번째로 작은 수를 푸는 문제로 바뀌게 된다. (C(A)는 A집합의 크기를 의미한다.)

 

이런식으로 재귀적인 접근법을 통해서 궁극적으로 k번째로 작은 수를 찾을 수 있게 되며, 퀵정렬 알고리즘의 특성을 그대로 활용할 수 있기 때문에 평균시간복잡도는 O(n) 최악공간복잡도는 O(n)에 검색 알고리즘을 풀 수 있게 된다.

 

퀵선택 알고리즘의 시각화. 출처: https://en.wikipedia.org/wiki/Quickselect

하지만, 여전히 축을 어떤식으로 고르느냐에 따라서 최악의 경우 시간복잡도가 O(n^2)으로 늘어날 수 있다. 이를 피할 수 있는 좋은 방법이 있을까?

중앙값들의 중앙값 (median-of-medians)

"중앙값들의 중앙값" 방법은 퀵선택 알고리즘의 축을 잘 설정하는 방법을 제시한다. 잘 설정한다는 것은 해당 방식을 도입해서 축을 설정한 퀵선택 알고리즘이 최악의 경우에도 선형 시간 복잡도를 가질 수 있도록 하는 것을 의미한다. 중앙값들의 중앙값 방법은 전체 배열을 크기가 5인 집합들로 나눈 뒤, 각 집합의 중앙값들을 구하고 그 중앙값들의 중앙값을 퀵정렬에 사용할 축으로 활용하는 것이다. 이제 (1) 이것을 어떻게 구현하는지와 (2) 이렇게 축을 설정하는게 왜 최악의 경우에도 퀵선택의 선형 시간 복잡도를 보장할 수 있는지에 대한 분석을 하는 것이 남았다.

구현

구현하기 위해 알아야 하는 간단한 사실은 중앙값을 구하는 문제가 검색 문제의 특수한 경우를 의미한다는 것이다. 중앙값은 n/2번째로 작은 수를 찾는 문제와 동일하다. (n이 배열의 크기이고 홀수일 때. 물론 짝수일 때에는 왼쪽 중앙값 혹은 오른쪽 중앙값을 사용해도 무방하다.) 즉, 우리가 크기가 5인 중앙값들의 중앙값을 찾는 것이 좀더 작은 배열 (정확히는 크기가 n/5로 줄어든 배열) 내에서 검색 문제를 푸는 것이므로, 분할 정복 알고리즘의 형태를 띈다는 것을 알 수 있다. 단순하게 앞에서부터 5개씩 배열을 자르고 그 배열 내의 중앙값은 상수 시간 내에 구할 수 있다고 가정하자. 그리고 중앙값에 해당하는 배열을 만든 뒤, 퀵선택 알고리즘을 재귀적으로 호출해서 중앙값을 찾고 그 값을 축으로 선택하도록 구성한다.

시간 복잡도 분석

해당 방식으로 축을 설정하는 것이 최악의 경우에도 선형 시간 복잡도를 제공한다는 것을 검증하기 위해, 점화식을 구성한다. 크기가 n인 배열을 위 알고리즘으로 실행시켰을 때 걸리는 시간을 T(n)이라고 설정하고 점화식을 구성해보면 다음과 같은 점화식이 만들어진다.

$$T(n)=c\cdot n + T\left(\frac{n}{5}\right) + T\left(n'\right)$$

첫번째 항은 주어진 배열을 크기가 5인 집합들로 나누고 각 집합 내의 중앙값을 상수 시간 내에 구하는 데 쓰이는 시간을 의미한다. 두번째 항은 구한 중앙값들의 중앙값을 쿽선택 알고리즘을 통해 구하는데 쓰이는 시간을 의미한다. 마지막항은 중앙값의 중앙값을 기반으로 두 집합 중에 하나를 선택하고 그 집합에 퀵선택 알고리즘을 통해 구하는데 쓰이는 시간을 의미한다. 여기서 중요한 것은 $n'$이 $\frac{7n}{10}$보다 반드시 작다는 점이다. 왜 그럴까?

 

그 이유는 중앙값들의 중앙값을 시각화해서 보면 알 수 있다. 아래 그림에서 각 열은 크기가 5인 집합을 의미 (열 내에서 정렬되어 있다.)하고 초록색은 각 집합 내의 중앙값을 의미한다. 보라색은 중앙값들의 중앙값을 의미하게 되는데 보라색을 기준으로 왼쪽 위에 해당하는 값들은 반드시 중앙값들의 중앙값보다 작을 수 밖에 없다. 마찬가지로, 오른쪽 아래에 해당하는 값들은 중앙값의 중앙값보다 클 수 밖에 없다. 따라서, 보라색 값 즉, 중앙값의 중앙값을 축으로 선택하면, 축을 기준으로 나누어진 두 집합 (축 값보다 작은 집합과 축 값보다 큰 집합)의 크기가 최소 $3 \cdot \frac{1}{2} \cdot \frac{n}{5}=\frac{3n}{10}$ 임을 보장할 수 있다. 이 말은, 어느 집합이든 크기가 $\frac{7n}{10}$을 넘을 수 없음을 의미한다.

 

중앙값들의 중앙값의 시각화. 출처: https://en.wikipedia.org/wiki/Selection_algorithm

그렇기 때문에 다음과 같은 부등식이 성립한다.

$$T(n) \le c \cdot n + T\left(\frac{n}{5}\right) + T\left(\frac{7n}{10}\right)$$

이는 정리하면 다음과 같은 식을 얻을 수 있으며, 최악의 경우에도 선형 시간 복잡도를 보장함을 의미한다.

$$T(n) \le 10 \cdot c \cdot n = O(n)$$

한계

이러한 방법이 최악의 경우에 대해서 시간 복잡도를 선형으로 낮춰주지만, 첫번째 항에 붙은 상수 시간이 꽤나 커 O(nlogn)의 시간복잡도를 가지는 알고리즘보다 실질적으로 느린 경우가 많다고 한다. 그리고, 사실 최악의 경우가 자주 발생하지 않는다는 사실 또한 실질적인 활용에 있어서 단순한 퀵정렬을 사용하는 게 더 좋을 수도 있음을 시사한다.

내재적인 분석을 통한 검색 알고리즘 (introselect)

앞서 살펴본 내용을 정리하면 다음과 같다. 퀵선택 알고리즘은 (장점) 평균 시간 복잡도를 선형으로 제공하지만, (단점) 최악의 경우 이차 시간 복잡도로 동작하게 된다. 이를 방지하기 위해 중앙값들의 중앙값으로 축을 선택하는 알고리즘을 도입하면 (장점) 최악의 경우에도 선형 시간 복잡도로 동작하도록 축을 고를 수 있지만, (단점) 중앙값들의 중앙값을 고르는데 많은 상수 시간을 쓰게 된다. 이 두 성질을 적절히 조합하여 사용하는 게 좋지 않을까?

이를 고려한 방법이 내재적인 분석을 통한 검색 알고리즘이며 방법은 단순히 퀵선택 알고리즘이 동작할 것이라고 낙천적으로 생각하고 k번 정도 작은 문제로 변형한 다음에, 작은 문제의 크기가 원래 문제의 크기의 절반보다 줄어들지 않을 경우, 축을 설정하는 방법으로 중앙값들의 중앙값을 설정하는 것을 의미한다. (다른 방법으로도 설정할 수 있다.) 해당 방법은 일반적인 경우에는 퀵정렬 알고리즘처럼 많은 상수 시간을 쓰지않고 O(n)에 동작한다는 장점이 있고, 일정 반복 횟수가 지난 뒤에는 중앙값들의 중앙값으로 축을 설정하기 때문에 퀵정렬 알고리즘의 최악의 경우 시간 복잡도인 O(n^2)보다 작은 O(n)에 동작하게 만들 수 있다.

numpy에서 적용

이 모든 구현이 numpy에 이미 구현되어 있으며, 해당 알고리즘을 쓰고 싶은 사람은 그냥 np.argpartition을 사용하면 된다. 이 글을 쓰게 된 계기가 이 api의 시간복잡도가 O(n)이라는 것에 놀라서 쓰게 되었다.

numpy.argpartition — NumPy v1.26 Manual

참고자료 (reference)

https://en.wikipedia.org/wiki/Selection_algorithm

https://en.wikipedia.org/wiki/Quickselect

https://en.wikipedia.org/wiki/Median_of_medians

https://en.wikipedia.org/wiki/Introselect

https://stackoverflow.com/questions/6910641/how-do-i-get-indices-of-n-maximum-values-in-a-numpy-array