[백준/2696] 중앙값 구하기 (Python)

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

2696번: 중앙값 구하기

 

Difficulty : Gold 2

 

Status : Solved

 

Time : 00:14:46

 

---

 

문제 설명

 

 

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

예를 들어, 수열이 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개씩 출력해야 한다.

 

입력 예시

 

3
9
1 2 3 4 5 6 7 8 9
9
9 8 7 6 5 4 3 2 1
23
23 41 13 22 -3 24 -31 -11 -8 -7
3 5 103 211 -311 -45 -67 -73 -81 -99
-33 24 56

 

출력 예시

 

5
1 2 3 4 5
5
9 8 7 6 5
12
23 23 22 22 13 3 5 5 3 -3
-7 -3

 

 

---

 

풀이

 

중앙값을 항상 업데이트할 수 있으려면

 

1. 항상 중앙값을 포함한 부분 집합이 정렬을 유지하고 있어야 하며

2. 값이 들어올 때마다 중앙값이 업데이트되는지 여부를 고려해야 한다.

3. 또한 중앙값을 탐색하는 데 걸리는 시간이 효율적이여야 한다.

 

따라서 입력이 들어올 때 정렬을 O(logN)만에 수행하며, 항상 최소/최댓값을 상수 시간 내에 접근 가능한 우선순위 큐를 고려해보자. 또한 중앙값을 기준으로 우선순위 큐 둘을 사용하는 것을 고려할 만하다. 중앙값 이하의 값들이 저장되는 최대 큐, 중앙값 초과 값들이 저장하는 최소 큐로 나누어 생각해 보자.

 

1. 입력이 들어올 때, 중앙값(즉 최대 큐의 루트 노드)과의 크기를 비교한다. 중앙값 이하라면 최대 큐에 삽입하고, 그렇지 않다면 최소 큐에 삽입한다.

2. 또한 중앙값의 특성상 최대 큐는 최소 큐와 크기가 같거나, 길이가 1 더 길어야 한다(이렇게 해야 최대 큐의 루트 노드가 항상 중앙값임이 보장된다) 만약 그렇지 않다면, 힙 삽입 및 삭제를 통해 크기를 조정해주도록 한다.

 

1, 2번 연산 모두 O(logN)내에 해결 가능하므로 주어진 시간 내에 충분히 해결 가능하다. 입출력의 조건만 유의해서 풀이해보도록 하자.

 

풀이 코드

from heapq import heappush, heappop
import sys
input = sys.stdin.readline

def solve() :
  M = int(input())
  num_list = list()
  for i in range((M-1) // 10 + 1) :
    num_list += list(map(int, input().split()))
  left_q, right_q, answer = list(), list(), list()
  left_q.append(-num_list[0])
  answer.append(num_list[0])
  
  for i in range(1, M) :
    if -left_q[0] < num_list[i] :
      heappush(right_q, num_list[i])
    else :
      heappush(left_q, -num_list[i])

    while len(left_q) - len(right_q) > 1 :
      heappush(right_q, -heappop(left_q))
    while len(left_q) - len(right_q) < 0 :
      heappush(left_q, -heappop(right_q))
    if not i % 2 :
      answer.append(-left_q[0])

  print(len(answer))
  for i in range(0, len(answer), 10) :
    print(*answer[i:min(i+10, len(answer))])

for _ in range(int(input())) :
  solve()

풀이 완료!