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

1. 문제 출처

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

1655번: 가운데를 말해요

2. 문제 풀이

그냥 막무가내로 풀었을 때는 정상적으로 답이 나왔으나 시간이 초과되었다.

n = int(input())

nums = []
nums.append(int(input()))
print(nums[0])

if n >=2:
    for i in range(1,n):
        temp = int(input())
        nums.append(temp)
        nums.sort()
        if len(nums)%2 == 0:
            middle_idx = int(len(nums)/2)
            print(min(nums[middle_idx-1] , nums[middle_idx]))
        else:
            print(nums[int((len(nums)-1)/2)])

n = int(input())
shouts_count = 1
nums = [10001] * n
nums[0] = int(input())
print(nums[0])

그래서 생각한 첫번째 해결책은 가운데 값을 기준으로 작은 값 큰 값 리스트들을 만들어서 sort를 나누어서 하였지만 역시나 안되었다.

n = int(input())

middle_list = [int(input()),int(input())]
print(middle_list[0])
middle_list.sort()
print(min(middle_list))
small_list = []
large_list = []

for i in range(2,n):
    temp = int(input())
    if (i+1) % 2 ==0:
        if temp<=middle_list[0]:
            if small_list[-1]<=temp:
                middle_list = [temp , middle_list[0]]
            else:
                middle_list = [small_list.pop() , middle_list[0]]
                small_list.append(temp)
        else:
            if large_list[-1]>=temp:
                middle_list = [ middle_list[0],temp]
            else:
                middle_list = [ middle_list[0],large_list.pop()]
                large_list.append(temp)
                
    else:
        if temp <middle_list[0]:
            small_list.append(temp)
            large_list.append(middle_list[1])
            
        elif middle_list[0]<=temp < middle_list[1]:
            small_list.append(middle_list[0])
            large_list.append(middle_list[1])
            middle_list[0] = temp
        else:
            large_list.append(temp)
            small_list.append(middle_list[0])
            middle_list[0] = middle_list[1]
    small_list.sort()
    large_list.sort(reverse=True)
    print(middle_list[0])

그래서 찬찬히 살펴보니 sort하는 것보다는 우선순위 큐로 Max Heap , Min Heap을 구현하여 정렬 하는 것이 O(logn) 의 시간 복잡도를 가지고 있는 것을 생각 하였다.

다만 파이썬의 내장 함수인 Heapq에는 max heap 이 구현되어 있지 않기 때문에 min heap에 값을 넣을 때 -1을 곱해서 넣는 방식으로 구현하였다.

다음은 중간을 기준으로 좌측은 max heap , 우측은 min Heap으로 구현한 코드이다.

import heapq
import sys
## input 보다 sys.stdin.readline() 더 빠르다.
## input 사용했을 때는 시간 초과 됨
n = int(sys.stdin.readline())

# 가운데값
middle = int(sys.stdin.readline())
print(middle)
# 각각 우측(min_heap) 좌측(max_heap) 이다
# 이름에 헷갈리지 말자.
min_heap = []
max_heap = []

# 돌아가는 것을 보고 싶으면 밑에 print문 주석을 풀어주자..
for i in range(1,n):
    # print("input:" , end = " ")
    temp = int(sys.stdin.readline())
    # 만약 입력값이 중앙값보다 크면
    if temp > middle:
        heapq.heappush(min_heap,temp)
        # 좌측 값들의 수가 우측값들의 수보다 2개 적으면 중앙값 갱신
        ## 중간이 2개일 때 큰값은 우측에 넣주기 때문에 2개이다.
        if len(min_heap) > (len(max_heap)+1):
            # max heap 이기 때문에 -1을 곱해서 넣어주어야 한다.
            heapq.heappush(max_heap , -middle)
            middle = heapq.heappop(min_heap)
            
    # 만약 입력값이 중앙값보다 작으면
    else:
        heapq.heappush(max_heap,-temp)
        # 우측값들의 수가 좌측값보다 작으면 중앙값 갱신
        if len(min_heap) < len(max_heap):
            heapq.heappush(min_heap , middle)
            middle = -(heapq.heappop(max_heap))
        
    # print("right:",min_heap)
    # print("left:",max_heap)
    # print("mid:",middle)
    print(middle)