콘텐츠로 이동

$O(N)$으로 $K$번째 순서 통계량 구하기

크기가 $N$인 배열 $A$와 수 $K$가 주어집니다. 이 문제는 배열에서 $K$번째로 큰 수, 즉 $K$번째 순서 통계량을 찾는 것입니다.

기본 아이디어는 퀵 정렬 알고리즘의 아이디어를 사용하는 것입니다. 사실 알고리즘 자체는 간단하지만, 퀵 정렬과 달리 이 알고리즘이 평균 $O(N)$에 동작함을 증명하기는 더 어렵습니다.

구현 (비재귀)

template <class T>
T order_statistics (std::vector<T> a, unsigned n, unsigned k)
{
    using std::swap;
    for (unsigned l=1, r=n; ; )
    {
        if (r <= l+1)
        {
            // the current part size is either 1 or 2, so it is easy to find the answer
            if (r == l+1 && a[r] < a[l])
                swap (a[l], a[r]);
            return a[k];
        }

        // ordering a[l], a[l+1], a[r]
        unsigned mid = (l + r) >> 1;
        swap (a[mid], a[l+1]);
        if (a[l] > a[r])
            swap (a[l], a[r]);
        if (a[l+1] > a[r])
            swap (a[l+1], a[r]);
        if (a[l] > a[l+1])
            swap (a[l], a[l+1]);

        // performing division
        // barrier is a[l + 1], i.e. median among a[l], a[l + 1], a[r]
        unsigned
            i = l+1,
            j = r;
        const T
            cur = a[l+1];
        for (;;)
        {
            while (a[++i] < cur) ;
            while (a[--j] > cur) ;
            if (i > j)
                break;
            swap (a[i], a[j]);
        }

        // inserting the barrier
        a[l+1] = a[j];
        a[j] = cur;

        // we continue to work in that part, which must contain the required element
        if (j >= k)
            r = j-1;
        if (j <= k)
            l = i;
    }
}

참고

  • 위의 무작위 알고리즘은 quickselect라고 부릅니다. 이 알고리즘이 올바르게 동작하려면 호출하기 전에 $A$를 무작위로 섞거나, 무작위 원소를 기준값으로 사용해야 합니다. median of medians와 같이 이 문제를 선형 시간에 해결하는 결정적 알고리즘도 있습니다.
  • std::nth_element는 C++에서 이 문제를 해결하지만, gcc의 구현은 최악의 경우 $O(n \log n )$ 시간에 동작합니다.
  • 가장 작은 원소 $K$개를 찾는 문제는 선형 추가 비용으로 $K$번째 원소를 찾는 문제로 환원할 수 있습니다. 그 원소들이 바로 $K$번째 원소보다 작은 원소들이기 때문입니다.

연습 문제