정렬하지 않고 n번째 원소를 뽑을 수 있어? - Quick Selection

1. Introduction

코딩을 하다 보면 꼭 전체를 정렬할 필요는 없는데 n번째로 큰 원소, 정확히는 n째 인덱스의 값을 얻고 싶을 때가 있다. 예를 들어, 한번 뽑고 나면 다시 섞여버리는 리스트라던가...

필자는 백준에서 문제를 풀다가 어떤 문제에서 존재를 알게되었는데, 처음에는 단순히 성능 "좋은 소팅 알고리즘을 요구하는 문제인가 보다.." 하고 생각했다. 하지만 계속 틀려서 다른 사람들의 답을 찾아보니 $O(n)$ 수준의 시간복잡도를 요구하는 문제였다. 그래서 이번엔 "선택 정렬에서 k까지만 정렬하고 바로 출력하자" 하고 제출하니 또 fail.

siba

생각해보니 이렇게 접근하면 k가 적을 때나 빠르지, 후반으로 갈수록 배열의 대부분을 정렬하니 걍 $O(n^2)$이나 다름 없었다.

그런데 더 찾아보니 k에 상관없이 항상 $O(n)$을 보장하는 알고리즘이 있었다.

---

2. Approach

Quick selection의 기본적인 아이디어는 Quick sort를 하되, Binary search마냥 k가 해당하는 위치만 계속 찾는(paritioning)하는 것이다.

이 방식을 이용하면 전체(N)를 한번 뒤져보고 (First partitioning) 그다음에 k가 포함된 parition(N/2)을 뒤져보고 (Second partitioning) ... 결국 partition의 크기가 1이 될 때까지 찾는다.

총 시간 복잡도가 N + N/2 + N/4 + ... 이고 이 값은 2N으로 수렴함을 우리는 잘 알고 있다.

따라서 Quick selection의 시간복잡도는 $O(n)$이다.

---

3. Evaluation

다음 코드는 내가 사용한 k까지의 선택정렬과 quick selection의 코드이다.

import random as rand
import time

def quick_selection(a, target):
    if len(a) == 0 or target < 0 or target >= len(a):
        return -1
    
    start = 0
    end = len(a) - 1
    
    while start < end :
        i = start
        j = end
        mid = a[(i+j)//2]
        
        while i < j:
            if a[i] >= mid:
                a[i], a[j] = a[j], a[i]
                j -= 1
            else:
                i += 1
            
        if a[i] > mid:
            i -= 1
        
        if target <= i:
            end = i
        else:
            start = i + 1
    
    return a[target]

def sort_selection(a, k):
    for i in range(0, k-1):

        min_idx = i

        for j in range(i + 1, k):
            if a[j] < a[min_idx]:
                min_idx = j

        a[i], a[min_idx] = a[min_idx], a[i]
        
    return a[k]

200000개의 배열에서 앞에 200번째 인덱스까지 찾는 경우의 결과는 아래와 같다.

다음 이미지는 2000번째 인덱스까지 늘린 경우다. 선택 정렬은 시간이 많이 늘었지만 quick selection은 그대로다.

 

---

4. Discussion

참고로 앞서 얘기했던 백준의 문제는 팬윅트리 (Fenwick Tree)라는 것으로도 풀 수 있다고 한다.