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

백준 파이썬 2696번 중앙값 구하기

Gold 2

 

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

import sys
input = sys.stdin.readline
for _ in range(int(input())):
    m = int(input())
    arr, q = [], []
    for _ in range(m//10+1):
        arr += list(map(int, input().split()))
    print((m+1)//2)
    for idx in range(m):
        q.append(arr[idx])
        q.sort()
        if idx%20 == 0 and idx!=0:
            print()
        if idx % 2 == 0:
            print(q[len(q)//2], end=" ")
    print()

중앙값 구하기로 heapq를 사용하지 않아도 정답 처리가 되기는 한다.

하지만 heapq를 사용해서 시간을 줄일 수 있다.

mid 의 값을 기준으로 왼쪽(작은 수)큐와 오른쪽(큰 수)큐를 생성하고 값을 빼준다.

import sys
import heapq
input = sys.stdin.readline
for _ in range(int(input())):
    m = int(input())
    arr = []
    for _ in range(m//10 + 1):
        arr += list(map(int, input().split()))
    l, r = [], []
    mid = arr[0]
    s = [mid]
    arr.pop(0)
    idx = 2
    for a in arr:
        if a < mid:
            heapq.heappush(l, -a) #mid와 가까운 수를 pop하기 위해 -값을 넣어준다.
        else:
            heapq.heappush(r, a)
        #print("left: ", l)
        #print("right: ", r)
        if idx % 2 == 1:
            if len(l) > len(r):
                heapq.heappush(r, mid)
                mid = -heapq.heappop(l)
            if len(l) < len(r):
                heapq.heappush(l, -mid)
                mid = heapq.heappop(r)
            s.append(mid)
        idx += 1
        #print("mid: ",mid)
    print(len(s))
    for i in range(len(s)):
        if i != 0 and i % 10 == 0:
            print()
        print(s[i], end=' ')
    print()

 

다음 예시에서 다음과 같은 출력값을 가지며, 시간 또한 줄어든 것을 확인할 수 있다.