[백준/우선순위큐,힙] 2696번. 중앙값 구하기(Python)

문제

중앙값 구하기

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

 

풀이

우리가 일반적으로 아는 방법은 일단, 입력을 다 받은 후 정렬하고 가운데 인덱스의 값을 출력하면 되지 않을까?이다.

현재 테스트 케이스 개수가 최대 1000개이고, 각 테스트 케이스 수열 크기가 약 10000개까지인데,
홀수번마다 정렬하여 가운데를 출력하는 방식을 사용하면,
길이가 n인 배열을 정렬할 때는 nlogN이 걸리고,
홀수번마다 정렬하므로 n/2번 x nlogN 해서 약 n^2 logN 으로 계산되고,
n이 약 10000(10^4)까지 가능하니까 
한개의 테스트 케이스에 10^8 log 10^4이고, 
최대 1000개의 테스트 케이스라면 위의 값에 1000을 곱한 10^3 x 10^8 log 10^4 이니까
1초(최대 가능 연산횟수 = 10^8) 안에는 해결할 수 없다.

따라서 이 방법은 사용할 수 없다.

 

그래서 우선순위큐를 사용한다.

현재 중앙값만 필요하고, 중앙값보다 큰 값들을 모은 pq(big heap)와 중앙값보다 작은 것들을 모은 pq(small heap)를 둔다.

big에서 가장 작은 값과 small에서 가장 큰 값과 지금 입력 받은 값, 총 3가지 값을 비교해서 그 중 중앙값을 고르면 그 값이 전체의 중앙값이다.

 

이제 코드로 구현하면된다.

 

파이썬에서 우선순위 큐 구현하기

파이썬에서는 힙(heap)기능을 위해 heapq라이브러리를 제공한다.

이 라이브러리는 우선순위 큐 기능을 구현하고자 할 때 사용된다.

heapq외에도 PriorityQueue라이브러리를 사용할 수 있지만,

코딩테스트 환경에서는 heapq가 더 빠르게 동작하기 때문에 heapq를 사용하는 것이 좋다.

 

파이썬의 힙은 최소힙(부모노드가 자식노드보다 크기가 작음)으로 구성되어 있다.

그리고 힙은 이진트리로 만들어져 있어, 삽입 삭제를 하는데 시간복잡도 O(NlogN)에 정렬이 완료된다.

 

import heapq

h = []

heapq.heappush(삽입할 리스트, 정수값이나 튜플/리스트)

튜플/리스트는 첫 번째 항목을 기준으로 우선순위를 가진다.

 

heaqp.heappop(삭제할 리스트)

pop할 때는 우선순위가 가장 높은 항목을 리턴하면서 삭제한다. (우선순위가 높다 = 값이 작다)

인자로는 리스트를 전달한다.

 

n개의 데이터를 삽입할 때 각 삽입마다 O(logN)이 걸린다. (트리이기 때문에 교환하면서 들어간다.)

따라서 총 O(nlogN)이 걸린다. 이 때 heapify를 이용하면 O(n)의 시간으로 배열을 heap으로 바꾼다.

 

최대힙의 경우,

튜플/리스트의 첫 번째 항목을 기준으로 우선순위를 가지기 때문에, 작은 것이 우선순위가 높은데,

1 -> 2 -> 3 에서 3 -> 2 -> 1로 바꾸고 싶으면 값에 -를 붙이면 된다.

 

 

다시 문제로 돌아와서,

중앙값을 비교할 때 small heap에서 가장 큰 값, big heap에서 가장 작은 값, 입력값 총 3개의 값을 비교해야 하기 때문에,

small heap은 max heap(최대힙)으로, big heap은 min heap(최소힙)으로 만든다.

코드는 아래와 같다.

# -*- coding: utf8 -*-
from sys import stdin
import heapq
# small은 최대힙 big은 최소힙으로 구성
# small에서 하나 뺀 거와 big에서 하나 뺀 거, 입력값1개 총 3개 중 중앙값을 구함.


def solution(data):
    smallh = []
    bigh = []
    middle = data[0]
    result = [middle]
    for idx, val in enumerate(data[1:], 1):
        if val > middle:
            heapq.heappush(bigh, val)
        else:
            heapq.heappush(smallh, (-val, val))

        if idx % 2 == 0:
            if len(smallh) < len(bigh):
                heapq.heappush(smallh, (-middle, middle))
                middle = heapq.heappop(bigh)
            elif len(smallh) > len(bigh):
                heapq.heappush(bigh, middle)
                middle = heapq.heappop(smallh)[1]
            result.append(middle)
    print(len(result))

    for i in range(len(result)):
        if i != 0 and (i+1) % 10 == 1:
            print()
        print(result[i], end=' ')
    print()


t = int(stdin.readline().rstrip())  # 테스트 케이스 개수
for i in range(t):
    m = int(stdin.readline().rstrip())  # 수열의 크기
    data = []
    if m % 10 == 0:
        for _ in range(m//10):
            data.extend(list(map(int, stdin.readline().rstrip().split(' '))))
    else:
        for _ in range(m//10+1):
            data.extend(list(map(int, stdin.readline().rstrip().split(' '))))

    solution(data)