Skip to main content Link Menu Expand (external link) Document Search Copy Copied

01655번 가운데를 말해요

Problem link: https://www.acmicpc.net/problem/1655

Keywords: Heap, Priority queue

부계를 만들어 문제를 다시 풀 때 느끼는 (서로 상반되는) 두 감정이 있다.

  • 다 까먹었는데, 새로 공부하게 되서 안심이 된다.
  • 그냥 머리에 답이 들어 있어 무지성으로 접근하게 되서 불안하다.

이 문제는 전형적으로 후자이다… 그냥 보자마자 min/max heap을 쓰는 문제라는 것을 알았기 때문이다.

여튼 풀이는 이러하다.

  • 가운데 값보다 작거나 같은 수들을 저장하고 있는 max heap, LEFT를 준비한다.
  • 가운데 값보다 큰 수들을 저장하고 있는 max heap, RIGHT를 준비한다.
  • 매번의 숫자 입력마다 아래를 수행한다.
    • Step 1:
      • if LEFT가 비어있다면: LEFT에 입력된 숫자를 넣는다.
      • else if LEFT.top() >= n: LEFT의 top을 RIGHT로 옮겨준다. 그리고, LEFT에 입력된 숫자를 넣는다.
      • else: RIGHT에 입력된 숫자를 넣는다.
    • Step 2:
      • Rebalancing 한다.
      • 즉, LEFT의 size보다 RIGHT의 size가 크다면, RIGHT의 top을 LEFT로 옮겨준다.

아래는 솔루션이다.

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

int N;
std::priority_queue<int, std::vector<int>> LEFT;   // max heap.
std::priority_queue<int, std::vector<int>> RIGHT;  // min heap.

int main() {
  // For faster IO.
  std::ios_base::sync_with_stdio(false);
  std::cout.tie(nullptr);
  std::cin.tie(nullptr);

  // Read input & solve.
  std::cin >> N;
  for (int i{0}; i < N; ++i) {
    int n;
    std::cin >> n;
    if (LEFT.size() == 0) {
      LEFT.push(n);
    } else if (n <= LEFT.top()) {
      RIGHT.push(-LEFT.top());
      LEFT.pop();
      LEFT.push(n);
    } else {
      RIGHT.push(-n);
    }
    // Rebalance.
    if (RIGHT.size() > LEFT.size()) {
      LEFT.push(-RIGHT.top());
      RIGHT.pop();
    }

    std::cout << LEFT.top() << "\n";
  }

  return 0;
}