[백준] 2696번: 중앙값 구하기

문제

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

2696번: 중앙값 구하기

 

풀이

두 개의 힙을 사용해 풀었습니다.

left는 최대힙, right는 최소힙입니다.

중앙값을 찾는 과정은 아래와 같습니다.

No.	input	left	mid	right
1	10	[1, 2, 3, 4]	5	[6, 7, 8, 9]
2		[1, 2, 3, 4]	5	[6, 7, 8, 9, 10]
3		[1, 2, 3, 4]	5	[6, 7, 8, 9, 10]
4	11	[1, 2, 3, 4]	5	[6, 7, 8, 9, 10]
5		[1, 2, 3, 4]	5	[6, 7, 8, 9, 10, 11]
6		[1, 2, 3, 4, 5]	6	[7, 8, 9, 10, 11]

1. 1~9까지의 수가 있고, 10이 입력값으로 들어옵니다.
2. 5 < 10이므로 right에 10이 들어갑니다.
3. 짝수 번째 입력이므로 다음으로 넘어갑니다.
4. 11이 입력값으로 들어옵니다.
5. 5 < 11이므로 right에 11이 들어갑니다.
6. 홀수 번째 입력이고 left의 크기보다 right의 크기가 더 크므로 현재 mid 값을 left에 넣고, mid 값을 right의 최솟값을 꺼내어 설정합니다.

 

코드

파이썬

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

def solution(nums):
    result = [nums[0]]
    mid = nums[0]
    left, right = [], []
    for i in range(1, M):
        if nums[i] > mid:
            heappush(right, nums[i])
        else:
            heappush(left, -nums[i])
        if i % 2 == 0:
            if len(left) < len(right):
                heappush(left, -mid)
                mid = heappop(right)
            elif len(left) > len(right):
                heappush(right, mid)
                mid = -heappop(left)
            result.append(mid)
    print(len(result))
    for i in range(len(result)):
        if i > 0 and i % 10 == 0:
            print()
        print(result[i], end=' ')
    print()

T = int(input())

for _ in range(T):
    M = int(input())
    nums = []
    if M % 10 == 0:
        for _ in range(M // 10):
            nums += list(map(int, input().split()))
    else:
        for _ in range(M // 10 + 1):
            nums += list(map(int, input().split()))
    solution(nums)