[Baekjoon]백준 2696 중앙값 구하기(골드 2) - Python

문제설명

문제를 살펴보면 M개의 수를 입력받아 홀수번째 index일때의 중앙값을 출력해야한다

당연히 이전문제처럼 가장 간단한 방법은 모든 값을 저장하여 index에 따라 값을 찾는 것이겠지만, 

정해진 시간과 메모리 안에서 해결하기 위해서는 다른 방법을 사용하여야한다

 

처음에는 하나의 힙으로 해결이 되는줄 알고 작성하다보니 중요한 점이 있었다

중간값을 출력해야하기때문에 하나의 힙으로는 조금 구현이 힘들것같아서 최대힙과 최소힙을 구현하였다

중간값은 최대힙에 저장하되, 음수로 저장하여 중앙값이 저장되게 하였고 최소힙에는 중간값보다 큰 값을 저장한다

이렇게 되면 최대힙에는 중간값보다 작은 값들이 음수로 큰 값부터 저장될것이고,(Ex -5, -3, -1)

최소힙에는 중간값보다 큰 값들이 저장될것이다(Ex 7, 8, 9)

대신 값을 그냥 넣으면 안 되기 때문에 최대힙과 최소힙의 크기를 일정 범위 이내로 유지하면서 값을 넣어주어야한다

만약 두 힙의 차이가 2이상이 되버리면 중간값이 아니게 되기 때문이다

if len(left) == len(right): # 힙 삽입 (크기 맞춰가며)
    heapq.heappush(left, -num)
else:
    heapq.heappush(right, num)

허나, 해당 방법을 사용하게 되면 중간값보다 값이 큰 경우에도 left(최대힙)에 값이 저장될 가능성이 존재한다

따라서 조건을 확인하여 swap하는 코드도 추가적으로 필요하다

# 중간값이 left[0]에 위치하기 위해 swap
if right and -left[0] > right[0]:
    l = -heapq.heappop(left)
    r = heapq.heappop(right)
    heapq.heappush(left, -r)
    heapq.heappush(right, l)

이후 중앙값을 저장하는 배열에 현재 돌고있는 index가 홀수인지 확인하여(해당 코드에서는 cnt) 홀수라면

left[0]에 저장된 값의 음수값을 리스트에 추가해준다

해당 for문이 종료될때마다(M개의 값을 입력받는 동작) result배열의 길이를 출력하고 해당 배열을 최대 10개씩

출력하면 된다 조금 복잡하긴 하지만 힙을 두 개 만든다면 해결가능한 문제이다

import sys, heapq
input = sys.stdin.readline

test_case = int(input())
for _ in range(test_case):
    size = int(input())
    result = [] # 중간값 저장
    left = [] # 중간값 이하의 숫자 저장(음수로 저장)
    right = [] # 중간값 초과의 숫자 저장
    cnt = 0
    for _ in range(size//10+1): # 10개씩 나누어 입력받기(문제조건)
        seq = list(map(int, input().split()))
        #print(seq)
        for num in seq:
            if cnt >= size:
                break

            if len(left) == len(right): # 힙 삽입 (크기 맞춰가며)
                heapq.heappush(left, -num)
            else:
                heapq.heappush(right, num)

            # 중간값이 left[0]에 위치하기 위해 swap
            if right and -left[0] > right[0]:
                l = -heapq.heappop(left)
                r = heapq.heappop(right)
                heapq.heappush(left, -r)
                heapq.heappush(right, l)

            cnt += 1
            if cnt % 2 == 1:  # 홀수 번째라면 중간값 저장
                result.append(-left[0])
        
    print(len(result))
    for num in range(0, len(result), 10):
        print(*result[num:num+10]) # 10개씩 출력