C++ 최대 힙(Max-Heap)에서 K번째로 큰 요소 찾기

이 튜토리얼에서는 최대 힙(max-heap)에서 k번째로 큰 요소를 찾는 프로그램을 작성해 보겠습니다.

힙 전체를 정렬하지 않고도 문제를 해결할 수 있는 효율적인 방법은 우선순위 큐(priority queue)를 활용하는 것입니다. 프로그램을 완성하기 위한 단계를 살펴보겠습니다.

문제 해결 접근 방법

• 올바른 값으로 최대 힙을 초기화합니다.
• 우선순위 큐를 생성하고 최대 힙의 루트 노드(가장 큰 값)를 삽입합니다.
• k - 1번 반복하는 루프를 작성합니다.
  • 큐에서 가장 큰 요소를 꺼냅니다(pop).
  • 꺼낸 노드의 왼쪽 자식과 오른쪽 자식을 우선순위 큐에 추가합니다.
• 루프가 끝난 후 우선순위 큐의 최상단에 있는 요소가 바로 k번째로 큰 요소입니다.
• 해당 값을 반환합니다.

이 방식은 매 반복마다 현재까지 확인된 후보 중 가장 큰 값만 추적하므로, 힙 전체를 정렬하는 것보다 훨씬 효율적입니다. 시간 복잡도는 O(k log k), 공간 복잡도는 O(k)입니다.

예제 코드

C++로 구현한 전체 코드를 살펴보겠습니다.

#include <bits/stdc++.h>
using namespace std;

struct Heap {
    vector<int> elements;
    int n;
    Heap(int i = 0): n(i) {
        elements = vector<int>(n);
    }
};

inline int leftIndex(int i) {
    return 2 * i + 1;
}

inline int rightIndex(int i) {
    return 2 * i + 2;
}

int findKthGreatestElement(Heap &heap, int k) {
    // (값, 인덱스) 쌍을 저장하는 우선순위 큐
    priority_queue<pair<int, int>> queue;
    queue.push(make_pair(heap.elements[0], 0));

    for (int i = 0; i < k - 1; ++i) {
        int node = queue.top().second;
        queue.pop();

        int left = leftIndex(node), right = rightIndex(node);

        if (left < heap.n) {
            queue.push(make_pair(heap.elements[left], left));
        }
        if (right < heap.n) {
            queue.push(make_pair(heap.elements[right], right));
        }
    }
    return queue.top().first;
}

int main() {
    Heap heap(10);
    heap.elements = vector<int>{ 44, 42, 35, 33, 31, 19, 27, 10, 26, 14 };
    cout << findKthGreatestElement(heap, 4) << endl;
    return 0;
}

실행 결과

위 코드를 실행하면 다음과 같은 결과를 얻을 수 있습니다.

33

예제 힙에서 가장 큰 요소부터 순서대로 44, 42, 35, 33이므로, 4번째로 큰 요소인 33이 정상적으로 출력됩니다.

마무리

이처럼 우선순위 큐를 활용하면 최대 힙에서 k번째로 큰 요소를 정렬 없이 간결하고 효율적으로 찾을 수 있습니다. 튜토리얼에 대해 궁금한 점이 있다면 댓글로 남겨주세요.