[백준] 1655. 가운데를 말해요

[백준] 1655. 가운데를 말해요 - 파이썬

 

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

 

풀이 1

일단 첫 시도는, 최소 힙을 사용.

입력받은 수들의 리스트를 저장 후, 반복문을 돌릴 때 마다 리스트를 알맞게 slicing하여 heapify,

그리고 매 반복마다 pop 후 heapify를 2/n회(가운데 값을 찾을 때 까지) 반복하여 중간값을 찾도록 코드를 짰다.

 

결과는 시간초과.

python에서의 heapify() 는 O(N)의 time complexity를 가지고,

heappop(), heappush() 는 각각 O(logN)의 time complexity를 가지므로,

간단히 계산해 봐도 최소 O(N^2) 이상의 complexity이고, 그냥 생각나는대로 대충 시도해봤기 때문에.. 아마 시간초과가 날 것이라고 예상은 했음.

import heapq

cnt = int(input())
inputList = []

for _ in range(cnt):
    inputList.append(int(input()))

print(inputList)

for i in range(cnt):
    # 매 반복마다 알맞게 리스트 slicing
    hq = inputList[:i+1]
    heapq.heapify(hq)   # 자른 리스트 heapify

    # 중간값을 찾을 때 까지 heappop (시간복잡도 약 O((n*logn), n = i+1)
    for _ in range(int((i+2)/2) - 1):
        heapq.heappop(hq)
    print(heapq.heappop(hq))

 

 

풀이 2

일단 최소 힙(Priority Queue, min Heap)을 쓰는 것 자체는 맞는 접근인 것 같고,

어떻게 해야 중간값을 빠르게 찾을 수 있을까 하는 고민을 굉장히 오랫동안 했다.

Heap 형태인 Priority Queue에서 index를 이용하여 중간값을 찾을 수는 없었기에..

 

그렇게 생각을 하다, 중간값을 바로바로 뽑아낼 수 있도록 저장 공간을 두 개로 나눠서, 중간값보다 큰 값과 작은 값을 넣는 곳을 따로 만들면 되지 않을까 생각을 했다.

그런데 중간값보다 큰 값과 중간값보다 작은 값의 그룹들에 각각 들어간 개수에 따라 현재의 중간값과 새로 들어오는 값을 비교하여 어디에 넣을 지 결정해 주어야 했고, 그 과정에서 '중간값보다 큰 그룹의 최솟값' 과 '중간값보다 작은 그룹의 최댓값' 을 확인해야 할 필요가 있었기에, 전자는 Min Heap, 후자는 Max Heap으로 처리하면 될 것이라고 생각했고, 그대로 코드를 짜 보았다.

 

구현하다 보니, middle이 중간값이려면 Min Heap의 개수가 Max Heap보다 1개 많거나, 같아야 한다는 것을 확인하고, 이에 대한 처리도 추가했다.

성공.

import heapq

cnt = int(input())
inputList = []

for _ in range(cnt):
    inputList.append(int(input()))

minHeap = []
maxHeap = []
middle = inputList[0]

for num in inputList[1:]:
    print(middle)

    if num > middle:
        heapq.heappush(minHeap, num)
    else:
        heapq.heappush(maxHeap, -num)  # Max-Heap을 사용하기 위해 음수로 삽입

    # minHeap(middle 이상의 그룹)이 maxHeap(middle 이하의 그룹)과 개수가 같거나 1개 많아야 함.
    if not (len(minHeap) == len(maxHeap) or len(minHeap) == len(maxHeap) + 1):
        # minHeap(middle 이상의 그룹)이 2개 더 많아져서 균형이 깨진 경우
        if (len(minHeap) == len(maxHeap) + 2):
            heapq.heappush(maxHeap, -middle)  # Max-Heap을 사용하기 위해 음수로 삽입
            middle = heapq.heappop(minHeap)
        # minHeap보다 maxHeap이 1개 더 많아져서 균형이 깨진 경우
        else:
            heapq.heappush(minHeap, middle)
            middle = -heapq.heappop(maxHeap)

print(middle)