[백준] 2696 중앙값 구하기

https://www.acmicpc.net/problem/2696

2696번: 중앙값 구하기

 

이번글은 백준사이트의 중앙값 구하기에 대한 설명입니다.

먼저 평균값과 중앙값은 다릅니다.

1 2 3 4 5가 있을 때 평균은 1+2+3+4+5 / 5 가 평균값이고 중앙값은 말 그대로 1~5사이에 있는 딱 중앙에 있는 값을 의미합니다.

 

그리고 우선 시간제한이 1초라는점과 입력값의 최대치를 확인해줍니다.

TC가 1000이고 하나의 TC에 대한 M의 입력이 9999입니다.

해당 문제를 입력받을 때마다 정렬하여 중앙값을 계산 할 수 도 있지만 그렇게 하게되면 

std::sort의 시간복잡도는 O(N LogN) 이기 때문에 시간 초과가 날 것입니다.(퀵 sort)

그렇다면 O(LogN)으로 더 줄여야 하기때문에 생각을 해보면 Heap을 사용하면 됩니다.

max heap, min heap을 통해서 중앙값을 구하는 것이죠.

 

최대 힙, 최소 힙에 데이터를 넣어주면서 중앙값을 구하는 것입니다.

처음에는 먼저 최대힙에 입력받은 데이터를 넣고 이 최대힙에 들어간 중앙값을 기준으로 최대 힙에 데이터를 넣을 지 최소힙에 데이터를 넣을 지 고민해주시면됩니다.

 

아래의 그림을 보시면 이해가 조금? 되실 껍니다 .위에 있는 힙이 최소힙, 아래가 최대 힙입니다.

주의할 점은 '균형'을 계속 맞춰 주어야 한다는 점입니다.

처음에 데이터가 5가 들어왔다면 이 5를 중앙값으로 잡고 그다음에 올 데이터가 중앙값인 5보다 작다면 

min heap에 max heap에 들어있던 top데이터를 넣어주고 max heap -> pop(), 이후 max heap에 중앙값 보다 작은 데이터를 넣어주어서 min heap, max heap의 데이터 개수의 균형을 맞춰 줍니다.

 

코드

#include <iostream>
using namespace std;

#include <list>
#include <vector>
#include <queue>

using ll = long long;
int loop, nums, num;

ll get(priority_queue<ll>* max, priority_queue<ll, vector<ll>, greater<ll>>* min)
{
    if (max->size() == min->size()) return (max->top() + min->top()) / 2.0;
    if (max->size() < min->size()) return min->top();
    return max->top();
}

int main()
{
    ios_base::sync_with_stdio(NULL);
    cin.tie(NULL);
    cout.tie(NULL);

    priority_queue<ll> maxHeap;
    priority_queue<ll, vector<ll>, greater<ll>> minHeap;
    vector<ll> vec;

    cin >> loop;
    for (int i = 0; i < loop; ++i)
    {
        cin >> nums;
        for (int j = 0; j < nums; ++j)
        {
            cin >> num;

            // TODO
            if (maxHeap.size() == 0)
            {
                maxHeap.push(num);

                if (j % 2 == 0)
                {
                    vec.push_back(get(&maxHeap, &minHeap));
                }

                continue;
            }

            if (maxHeap.size() == minHeap.size())
            {
                if (num < get(&maxHeap, &minHeap)) maxHeap.push(num);
                else minHeap.push(num);

                if (j % 2 == 0)
                {
                    vec.push_back(get(&maxHeap, &minHeap));
                }

                continue;
            }

            if (maxHeap.size() < minHeap.size())
            {
                if (num > get(&maxHeap, &minHeap))
                {
                    maxHeap.push(minHeap.top());
                    minHeap.pop();
                    minHeap.push(num);
                }
                else maxHeap.push(num);

                if (j % 2 == 0)
                {
                    vec.push_back(get(&maxHeap, &minHeap));
                }

                continue;
            }

            if (num < get(&maxHeap, &minHeap))
            {
                minHeap.push(maxHeap.top());
                maxHeap.pop();
                maxHeap.push(num);
            }
            else minHeap.push(num);

            if (j % 2 == 0)
            {
                vec.push_back(get(&maxHeap, &minHeap));
            }
        }

        if (nums % 2 == 0) cout << nums / 2 << endl;
        else cout << nums / 2 + 1 << endl;

        for (ll v : vec)
        {
            cout << v << " ";
        }
        cout << endl;

        while (maxHeap.empty() == false) maxHeap.pop();
        while (minHeap.empty() == false) minHeap.pop();
        vec.clear();
    }

    return 0;
}