[백준] 2696번 중앙값 구하기 - C++/Java

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

2696번: 중앙값 구하기

문제

어떤 수열을 읽고, 홀수번째 수를 읽을 때 마다, 지금까지 입력받은 값의 중앙값을 출력하는 프로그램을 작성하시오.

예를 들어, 수열이 1, 5, 4, 3, 2 이면, 홀수번째 수는 1번째 수, 3번째 수, 5번째 수이고, 1번째 수를 읽었을 때 중앙값은 1, 3번째 수를 읽었을 때는 4, 5번째 수를 읽었을 때는 3이다.

입력

첫째 줄에 테스트 케이스의 개수 T(1 ≤ T ≤ 1,000)가 주어진다. 각 테스트 케이스의 첫째 줄에는 수열의 크기 M(1 ≤ M ≤ 9999, M은 홀수)이 주어지고, 그 다음 줄부터 이 수열의 원소가 차례대로 주어진다. 원소는 한 줄에 10개씩 나누어져있고, 32비트 부호있는 정수이다.

출력

각 테스트 케이스에 대해 첫째 줄에 출력하는 중앙값의 개수를 출력하고, 둘째 줄에는 홀수 번째 수를 읽을 때 마다 구한 중앙값을 차례대로 공백으로 구분하여 출력한다. 이때, 한 줄에 10개씩 출력해야 한다.

---

    테스트 케이스가 최대 1000개 이기 때문에 1초에 1억번 계산을 한다고 가정했을 때 한 가지 테스트 케이스를 10만 번의 계산 안에 풀어야 한다. 따라서 O(N^2)으로 풀 수는 없고 O(NlogN)으로 풀어야 한다. 이 말인즉슨 정렬이나 배열을 순회하며 중앙값을 구할 수는 없다는 뜻이다.

    따라서 대표적인 삽입에 O(NlogN)의 시간을 가지는 우선순위큐를 활용하였다. 중앙값인 pivot을 두고 pivot 보다 작은 값들의 PQ (lpq)와 큰 값들의 PQ(rpq)를 둔다. 이때 lpq는 숫자가 큰 것이 top() , rpq는 작은 것이 top() 일 수 있게 rpq는 greater 옵션을 준다. 첫 숫자는 pivot으로 입력을 받고 바로 중앙값에 저장을 한다. 그다음 숫자 두 개 씩 입력받으며 적절한 PQ에 삽입을 한다. 삽입을 완료한 후 두 PQ의 크기가 같다면 중앙값은 유지되고 크기가 다르다면 pivot을 크기가 작은 쪽 PQ에 삽입하고 새로운 pivot은 크기가 큰 쪽 PQ의 top()으로 바꾸게 된다면 두 PQ의 크기가 일치하여 pivot이 중앙값이 된다.

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

int tc;

int main() {
	cin >> tc;
	while (tc--) {
		vector<int> mid;
		int n,pivot,a,b;
		
		priority_queue <int, vector<int>, greater<int>> rpq;
		priority_queue<int> lpq;

		cin >> n >> pivot; mid.push_back(pivot);
		for (int i = 2; i <= n; i+=2) {
			cin >> a >> b;
			if (a >= pivot)	rpq.push(a);
			else lpq.push(a);

			if (b >= pivot) rpq.push(b);
			else lpq.push(b);
			
			if (lpq.size() != rpq.size()) {
				if (lpq.size() > rpq.size()) {
					rpq.push(pivot);
					pivot = lpq.top();
					lpq.pop();
				}
				else {
					lpq.push(pivot);
					pivot = rpq.top();
					rpq.pop();
				}
			}
			mid.push_back(pivot);			
		}

		cout << mid.size() << "\n";
		for (int i = 0; i < mid.size(); i++) {
			cout << mid[i] << " ";
			if (i % 10 == 9 ) cout << "\n";
		}

	}

}

 

import java.util.*;

public class Main {
    public static void main (String[] args){
        Scanner sc = new Scanner(System.in);
        Integer tc;
        tc = sc.nextInt();
        while(true){
            if(tc-- == 0) break;
            PriorityQueue<Integer> lpq = new PriorityQueue<>(Collections.reverseOrder());
            PriorityQueue<Integer> rpq = new PriorityQueue<>();
            ArrayList<Integer> mid = new ArrayList<>();
            Integer n,a,b,pivot;
            n = sc.nextInt();
            pivot = sc.nextInt();
            mid.add(pivot);

            for(int i=2;i<=n;i+=2){
                a = sc.nextInt();
                b = sc.nextInt();

                if(a >=pivot) rpq.add(a);
                else lpq.add(a);

                if(b >= pivot) rpq.add(b);
                else lpq.add(b);

                if(lpq.size() != rpq.size()){
                    if(lpq.size() > rpq.size()){
                        rpq.add(pivot);
                        pivot = lpq.poll();
                    }
                    else{
                        lpq.add(pivot);
                        pivot = rpq.poll();
                    }
                }
                mid.add(pivot);
            }
            System.out.println(mid.size());
            for (int i = 0; i < mid.size(); i++) {
                System.out.print(mid.get(i) + " ");
                if (i % 10 == 9 ) System.out.println();
            }
            System.out.println();
        }
    }
}

---

    위에서 설명한 것과 같은 사고의 흐름을 가지고 문제를 원활히 푼 것 같다. 시간 복잡도 계산부터 정석적으로 푼 것 같아서 만족스럽다. 언젠가 두개의 PQ를 이용하는 문제를 푼 것 같은데 그때는 어떻게 두 개를 쓰는 거지 싶었는데 이번에는 바로 감이 와 풀 수 있었다. 역시 문제를 많이 풀어두면 언젠가는 도움이 되나 보다.

     Java의 경우 ArrayList 나 PQ 모두 삽입을 할 때 add를 써서 편리한 것 같다. 처음으로 IntelliJ를 활용하여 Java를 풀었는데 처음이라 낯설다. 하지만 다들 이제는 eclipse 보다 IntelliJ를 사용한다고 하니 익숙해져야겠다!