[Heap] 백준 1655번 “가운데를 말해요” Python 풀이

문제 링크

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

풀이

• 우선, 최대힙을 위한 leftheap과 최소힙을 위한 rightheap을 생성합니다.
  • leftheap은 중간값보다 같거나 작은값으로, rightheap은 중간값보다 큰값으로 유지하기 위함입니다.
• 이후, 요소들을 순서대로 입력 받으며, 아래 알고리즘을 진행합니다.
  • 현재 leftheap과 righteap의 길이가 같다면, (현재 요소 * -1) 을 leftheap에 삽입하고, 길이가 같지 않다면, (현재 요소)를 rightheap에 삽입합니다.
    • 파이썬의 heapq는 기본적으로 최소 힙만을 제공하기 때문에,
    • -1 을 곱함으로써 leftheap을 최대힙으로 만들 수 있습니다.
  • 이후, leftheap의 루트와 rightheap의 루트를 비교합니다.
    • leftheap과 rightheap에 요소가 존재할 때,
    • leftheap의 (루트 * -1)이 rightheap의 (루트)보다 크다면,
    • leftheap의 (루트 * -1)과 rightheap의 (루트)를 바꿉니다.
      • 변환 시 음수, 양수에 대한 전처리도 진행해야합니다.
• 모든 요소들을 처리했다면, leftheap의 (루트 * -1)를 출력합니다.
• 입력 예제는 아래와 같습니다.
  • ex) [1, 5, 2, 10, -99, 7, 5]
  • num = 1 👉🏻 left = [-1] / right = []
  • num = 5 👉🏻 left = [-1], right = [5]
  • num = 2 👉🏻 left = [-2,-1], right = [5]
  • num = 10 👉🏻 left = [-2,-1], right = [5,10]
  • num = -99 👉🏻 left = [-2,-1,99], right = [5,10]
  • num = 7 👉🏻 left = [-2,-1,99], right = [5,7,10]
  • num = 5 👉🏻 left = [-5,-2,-1,99], right = [5,7,10]

코드

import sys
import heapq

n = int(sys.stdin.readline())

leftheap = []
rightheap = []

for _ in range(n):

    # 현재 숫자
    num = int(sys.stdin.readline())

    # leftheap의 길이가 rightheap의 길이와 같다면
    if (len(leftheap) == len(rightheap)):
        # leftheap에 현재 숫자를 음수 형태로 삽입
        heapq.heappush(leftheap, -num)

    # leftheap의 길이가 rightheap의 길이와 다르다면
    else:
        # rightheap에 현재 숫자를 삽입
        heapq.heappush(rightheap, num)

    # 양쪽 heap에 요소들이 존재한다면,
    # leftheap의 루트에 음수를 곱한 값이 rightheap의 루트보다 크다면
    if (len(leftheap) >= 1 and len(rightheap) >= 1 and -leftheap[0] > rightheap[0]):

        # leftheap의 루트를 제거하고, 
        leftroot = heapq.heappop(leftheap)
        # leftheap의 루트에 음수를 곱한 값을 rightheap에 삽입
        heapq.heappush(rightheap, -leftroot)

        # rightheap의 루트를 제거하고, 
        rightroot = heapq.heappop(rightheap)
        # rightheap의 루트에 음수를 곱한 값을 leftheap에 삽입
        heapq.heappush(leftheap, -rightroot)

    # leftheap의 루트에 음수를 곱한 뒤 출력
    print(-leftheap[0])

참고

• https://velog.io/@uoayop/BOJ-1655-가운데를-말해요Python