C++: Priority Queue 구현하기

코딩 테스트 문제를 풀면서 PQ(Priority Queue)를 직접 구현하여야 할 경우가 있었다. 단순히 Parent, Left Child, Right Child의 인덱스만 아는 것이 아니고, push, pop, updaet, erase 모든 기능에 대해 어떤 식으로 동작하는지 간략히 정리해보고자 한다.

PQ에 삽입할 데이터

struct Info {
    int idx;
    int score;
}

 PQ에 삽입할 데이터는 일반 변수가 아닌, 사용자 지정 변수를 기반으로 정렬이 가능한 경우를 다루고자 한다. 위 구조체 처럼 score를 기준으로 정렬될 수 도 있고, 경우에 따라 다양한 우선순위에 따라 정렬시킬 수 있다. `idx`의 경우, PQ에 push, pop을 할 때는 필요하지 않지만, PQ에 update, erase 할 경우에 필요하다.

 

PQ init

void init() {
    // size를 0으로 만들면 top부터 다시 삽입하게 된다.
    size = 0;
}

 PQ에 init는 간단하다. 현재 사용중인 heap의 크기를 0으로 만들어버리면 top부터 다시 값을 채우게 된다.

 

PQ Push

int updateParent(int idx) {
    int curIdx = idx;

    // current가 root가 아니고, parent 노드보다 우선순위이면
    while (curIdx > 0 && comp(curIdx, (curIdx - 1) / 2)) {
        swap(curIdx, (curIdx - 1) / 2);
        curIdx = (curIdx - 1) / 2;
    }

    return curIdx;
}

void push(Info* r) {
    heap[size] = r;
    r->idx = size;
    updateParent(size);
    size++;
}

 PQ에 새로운 값을 push한 경우, 탐색하고자 하는 index가 root가 아니고, parent보다 우선수위가 높을 때까지 계속 탐색하여, 현재 삽입한 값의 자리를 찾아야 한다. 이를 위해, `swap`, `comp` 함수를 정의하여 호출하면 보다 깔끔하게 작성할 수 있다. 여기서 중요한 것은 parent의 index는 `(current index - 1) / 2`라는 것이다.

 

PQ pop

void updateChild(int idx) {
    int curIdx = idx;
    // current가 왼쪽 자식이 있을 때 까지
    while (curIdx * 2 + 1 < size) {
    int child;

    if (curIdx * 2 + 2 == size)
        child = curIdx * 2 + 1;
    else
        child = comp(curIdx * 2 + 1, curIdx * 2 + 2) ? curIdx * 2 + 1 : curIdx * 2 + 2;

    if (comp(curIdx, child))
        break;

        swap(curIdx, child);
        curIdx = child;
    }
}

Info* pop() {
    if (size == 0)
        return nullptr;

    Info* top = heap[0];
    size--;
    swap(0, size);
    updateChild(0);
    return top;
}

 PQ의 top을 pop하는 경우에는 0번 인덱스의 값을 return 하고, 0번과 마지막 인덱스와 swap 한 뒤에, 해당 값의 위치를 찾으면서 다시 정렬하는 과정을 거쳐야 한다. 여기서 중요한 것은 left child는 `current index * 2 + 1`, right child는 `current index * 2 + 2`로 찾을 수 있다는 것이다.

 

PQ update

  PQ update는 heap의 특정 인덱스에 대한 값이 업데이트 되면, `updateChild(updateParent(idx))`와 같이 특정 인덱스에 대한 위치를 다시 정렬하는 방식으로 진행하면 된다.

 

PQ erase

 PQ erase는 pop은 0번 인덱스와 마지막 인덱스를 swap하였다면, 이와 달리 특정 인덱스와 마지막 인덱스를 swap 하는 차이만 가지고 있다. update와 마찬가지로 중간 인덱스에 erase가 발생한다면, child, parent 모두 업데이트하여야 하므로 `updateChild(updateParent(idx))`와 같이 호출하여 다시 정렬하는 방식으로 진행한다.

 

전체 코드

#define MAX_HEAP	100

struct Info {
    int idx;
    int score;
};

struct PQ {
    /*
    *  Parent = (idx - 1) / 2
    *  leftChild = idx * 2 + 1
    *  rightChild = idx * 2 + 2
    */
    int size;
    Info* heap[MAX_HEAP];

    void init() {
        // size를 0으로 만들면 top부터 다시 삽입하게 된다.
        size = 0;
    }

    bool comp(int a, int b) {
        // 부등호를 반대로하면 max heap이 된다.
        return heap[a]->score < heap[b]->score;
    }

    void swap(int a, int b) {
        heap[a]->idx = b;
        heap[b]->idx = a;
        Info* temp = heap[a];
        heap[a] = heap[b];
        heap[b] = temp;
    }

    int updateParent(int idx) {
        int curIdx = idx;

        // current가 root가 아니고, parent 노드보다 우선순위이면
        while (curIdx > 0 && comp(curIdx, (curIdx - 1) / 2)) {
            swap(curIdx, (curIdx - 1) / 2);
            curIdx = (curIdx - 1) / 2;
        }

        return curIdx;
    }

    void updateChild(int idx) {
        int curIdx = idx;
        // current가 왼쪽 자식이 있을 때 까지
        while (curIdx * 2 + 1 < size) {
            int child;

            if (curIdx * 2 + 2 == size)
                child = curIdx * 2 + 1;
            else
                child = comp(curIdx * 2 + 1, curIdx * 2 + 2) ? curIdx * 2 + 1 : curIdx * 2 + 2;

            if (comp(curIdx, child))
                break;

            swap(curIdx, child);
            curIdx = child;
        }
    }

    void push(Info* r) {
        heap[size] = r;
        r->idx = size;
        updateParent(size);
        size++;
    }

    Info* pop() {
        if (size == 0)
            return nullptr;

        Info* top = heap[0];
        size--;
        swap(0, size);
        updateChild(0);
        return top;
    }
    /*
    * update, erase는 특정부분에서 진행되므로
    * parent, child 모두 업데이트를 하여야 함.
    */
    void update(Info* r) {
        updateChild(updateParent(r->idx));
    }

    void erase(Info* r) {
        int idx = r->idx;
        size--;
        swap(idx, size);
        updateChild(updateParent(idx));
    }
};

 전체 코드는 위와 같다. 데이터의 삽입, 삭제, 갱신에 따라 parent, child를 업데이트할지 알고, 이를 탐색하기 위한 로직을 위해한다면 쉽게 작성할 수 있다. update, erase를 위해서는 함수를 잘 구분하여 구현하는 것도 중요한 것 같다.