백준 1655번 - 평범한 배낭

문제를 읽고 단순히 접근하면 입력받을 때 마다 정렬하여 중간값을 찾아 출력하면 되겠다고 생각합니다. 저 또한 그러했고, 쉽게 코딩 후 결과는 시간초과...

정렬과 중간값을 찾는 과정을 O(1)로 만들어야 한다는 것을 깨달은 후 여러 자료구조 형태를 고민하다 우선순위 큐를 사용해보았습니다.

이 문제는 우선순위 큐를 2개를 사용하여 입력되는 숫자들의 중간값을 top로 유지시키며 추출하는 문제입니다.

제가 처음 구상한 방법은 하나의 정렬된 상태를 만들기 위해 우선순위 큐 두개(max heap 1개, min heap 1개)로 나누는 방법입니다.

예를들어

입력된 숫자가 1, 2, 3, 4, 5, 6, 7, 8 이라고 하면

우선순위 큐1(max heap) : 4, 3, 2, 1
우선순위 큐2(min heap) : 5, 6, 7, 8

처럼 구성하는 방법입니다.

2개로 나뉘어져 있지만 우선순위 큐1을 거꾸로 시작하여 반시계방향으로 우선순위 큐2까지 보면 순서대로 1, 2, 3, 4, 5, 6, 7, 8임을 알 수 있고, 우선순위 큐1에 항상 중간값이 오는 걸 알 수 있습니다.

로직은 먼저 max heap에 숫자를 처음 입력받습니다.

그 수보다 크면 min heap, 작으면 max heap에 push합니다.

이후 max heap의 top를 기준으로 새로 입력받는 숫자가 작으면 max heap, 크면 mean heap으로 push합니다.

두 우선순위 큐의 갯수를 항상 같거나 1씩 차이나도록 설정해야 max heap에 중간값이 위치할 수 있습니다.

우선순위 큐1(max heap) : 2 1
우선순위 큐2(min heap) : 3 4 5 6 7 8

위 예시처럼 갯수가 한쪽으로 편향되면 안됩니다. 새로운 숫자를 입력받아 push를 할 때 두 큐의 개수가 같으면 한쪽에 push를 하고 나서 push한 큐의 top를 다른 큐에 옮깁니다.

두 큐의 개수가 다를 경우 push하기 전에 옮겨서 개수를 같게 맞춘 다음 위와 같은 작업을 실시합니다.

이렇게 두 우선순위 큐의 균형을 맞춰가며 중간값을 유지하는 로직으로 구현하였습니다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <queue>

using namespace std;


int main(){
    ios_base::sync_with_stdio(false); cin.tie(NULL);	
    int num;
    cin >> num;
    priority_queue<int, vector<int>, greater<int> > min_queue;
    priority_queue<int, vector<int>, less<int> > max_queue;

    int input;
    cin >> input;
    max_queue.push(input);
    cout << max_queue.top() << '\n';
    for(int i = 1; i < num; i++){
        
        cin >> input;
        
        

        if(max_queue.top() < input){
            if(max_queue.size() == min_queue.size()){
                min_queue.push(input);
                int temp = min_queue.top();
                min_queue.pop();
                max_queue.push(temp);
            }
            else{
                min_queue.push(input);
                if(max_queue.size() != min_queue.size()){
                    int temp = min_queue.top();
                    min_queue.pop();
                    max_queue.push(temp);
                }
            }
            
        }
        else{
            if(max_queue.size() == min_queue.size()){
                max_queue.push(input);
                int temp = max_queue.top();
                max_queue.pop();
                min_queue.push(temp);
            }
            else{
                max_queue.push(input);
                if(max_queue.size() != min_queue.size()){
                    int temp = max_queue.top();
                    max_queue.pop();
                    min_queue.push(temp);
                }
            }
            
        }
        if(max_queue.size() < min_queue.size()){
            cout << min_queue.top() << '\n';
        }
        else{
            cout << max_queue.top() << '\n';

        }
        
        
    }
    
}

다른 분의 코드 참조 내용

다른 분들의 코드를 보며 공부를 하다 저의 로직보다 간단한 로직이 있었습니다. 제가 구현할 때 push할 때 max heap의 top와 비교하여 우선순위 큐를 골라 push하였는데 이는 필요없는 과정이였습니다. 항상 두 우선순위 큐의 갯수만 같도록 max heap과 min heap의 개수가 같을 때 max heap에 push, 개수가 다를 때 min heap에 push를 통해 두 큐의 개수 유지를 합니다.(이 부분은 저와 같습니다.) 하지만 입력되는 수의 크기를 비교하여 push할 큐를 고르는 것이 아니라 push를 한 후 두 큐의 top을 비교하여 max heap의 top에 항상 작은 값들을 오게 합니다. 이는 우선순위 큐에 push 되는 것만으로 정렬이 되고, top를 비교하여 max heap에 작은 값을 옮겨준다면 최종형태가 저의 로직과 같아지는 것을 알게 되었습니다. 이 로직이 훨씬 조건문들이 적고 간편하며 직관적이라 생각됩니다.

#include <iostream>
#include <queue>
using namespace std;

int main(){
    ios_base::sync_with_stdio(false); cin.tie(NULL);	
    priority_queue<int , vector<int>, greater<int> > min_queue;
    priority_queue<int, vector<int>, less<int> > max_queue;

    int num;
    cin >> num;
    for(int i = 0; i < num; i++){
        int input;
        cin >> input;
        if(max_queue.size() == min_queue.size()){
            max_queue.push(input);
        }
        else{
            min_queue.push(input);
        }
        // 두 우선순위 큐가 비어있지 않으면서 max heap의 top가 min heap의 top 보다 클 경우 swap
        if(!max_queue.empty() && !min_queue.empty() && max_queue.top() > min_queue.top()){
            int swap_temp1 = max_queue.top();
            int swap_temp2 = min_queue.top();
            max_queue.pop();
            min_queue.pop();
            max_queue.push(swap_temp2);
            min_queue.push(swap_temp1);
        }
        
        cout << max_queue.top() << '\n';
    }
}

참조

https://blog.naver.com/music4537/222249608344

추가적으로

위 문제처럼 입출력이 n번만큼 자주 등장할 경우 endl이나 cin, cout을 사용하면 시간초과가 발생할 확률이 높습니다.

ios_base::sync_with_stdio(false); cin.tie(NULL);

위 처리를 메인 내에서 해준다면 입출력속도를 향상시킬 수 있습니다. 또한 줄바꿈을 'endl' 대신 '\n'을 사용하는 것이 더 빠릅니다.