[python] 백준 1655 :: 가운데를 말해요 (최대힙, 최소힙, 우선순위 큐)

[가운데를 말해요]

# 문제
백준이는 동생에게 "가운데를 말해요" 게임을 가르쳐주고 있다. 백준이가 정수를 하나씩 외칠때마다 동생은 지금까지 백준이가 말한 수 중에서 중간값을 말해야 한다. 만약, 그동안 백준이가 외친 수의 개수가 짝수개라면 중간에 있는 두 수 중에서 작은 수를 말해야 한다.

예를 들어 백준이가 동생에게 1, 5, 2, 10, -99, 7, 5를 순서대로 외쳤다고 하면, 동생은 1, 1, 2, 2, 2, 2, 5를 차례대로 말해야 한다. 백준이가 외치는 수가 주어졌을 때, 동생이 말해야 하는 수를 구하는 프로그램을 작성하시오.

# 입력
첫째 줄에는 백준이가 외치는 정수의 개수 N이 주어진다. N은 1보다 크거나 같고, 100,000보다 작거나 같은 자연수이다. 그 다음 N줄에 걸쳐서 백준이가 외치는 정수가 차례대로 주어진다. 정수는 -10,000보다 크거나 같고, 10,000보다 작거나 같다.

# 출력
한 줄에 하나씩 N줄에 걸쳐 백준이의 동생이 말해야 하는 수를 순서대로 출력한다.

풀이

전체 코드

import sys
import heapq

n = int(input())
left = []   # 중간값을 포함한 이전 값 최대힙
right = []  # 중간값 이후의 값 최소힙

for i in range(n):
    x = int(sys.stdin.readline())

    # left가 비어있거나
    # 현재 입력값이 left에서 제일 작은 값보다 작거나 같으면
    if len(left) == 0 or -left[0] >= x:
        heapq.heappush(left, -x)
    else:
        heapq.heappush(right, x)

    # left가 right에 1을 더한 것보다 더 길면
    if len(left) > len(right) + 1:
        # left에서 제일 큰 값을 빼서 right에 넣습니다.
        temp = -heapq.heappop(left)
        heapq.heappush(right, temp)
    # right가 left보다 길면
    elif len(right) > len(left):
        # right에서 제일 작은 값을 빼서 left에 넣습니다.
        temp = heapq.heappop(right)
        heapq.heappush(left, -temp)

    # 중간값은 최대힙에서 제일 큰 값입니다.
    print(-left[0])

 

0. 풀이 방법

• 최대힙과 최소힙은 각각 내림차순과 오름차순으로 데이터를 정렬해서 저장하는 자료구조이다.
• 이 문제처럼 정렬된 값에서 가장 큰 값, 혹은 가장 작은 값을 구해야 할 때 아주 유용하다. 너무 맘에 든당 😍
• 중간값을 기준으로, 들어온 값이 중간값보다 크거나 같은지를 확인해서 왼쪽 left 와 오른쪽 right 로 나눠서 담을 것이다.
• 값을 각 힙에 담아나가면서 left와 right의 길이가 같거나 최대 1만큼만 차이가 나도록 맞춰줘야 하는데,
  left에서는 제일 큰 값을, right에서는 제일 작은 값을 뽑아서(pop) 반대쪽에 넣어준다.
  (이러면 바로 정렬이 되어버림👍🏻)
• left에서는 제일 큰 값을 빼내야 하고 right에서는 제일 작은 값을 빼내야 하므로
  left는 최대힙, right는 최소힙으로 활용할 것이다.
• 그렇다면 중간값은 어디에 담는 게 좋을까?!
  수의 개수가 짝수일 때는 더 작은 값을 출력해야 하니 left에서 제일 큰 값을 출력해야 한다.
  따라서, 중간값을 left에 넣어놓으면
  수의 개수가 짝수일 때도 left에서, 홀수일 때도 left에서 꺼내면 되니까 편리하다.

 

1. 입력받은 값 힙에 넣기

• 우선 left와 right 리스트를 만든다.

left = []   # 중간값을 포함한 이전 값 최대힙
right = []  # 중간값 이후의 값 최소힙

• 하나씩 입력받은 값을 현재 중간값인 left에서 제일 큰 값 -left[0]와 비교해서 left에 넣을지 right에 넣을지 결정한다.
• 중간값보다 작거나 같으면 left에 넣는다.
• left가 비어있다면 left에 넣어버린다.

for i in range(n):
    x = int(sys.stdin.readline())

    # left가 비어있거나
    # 현재 입력값이 left에서 제일 큰 값보다 작거나 같으면
    if len(left) == 0 or -left[0] >= x:
        heapq.heappush(left, -x)
    else:
        heapq.heappush(right, x)

 

2. left와 right 길이 맞춰주기

• 값의 개수가 짝수일 때는 left와 right의 길이가 동일해야 하고,
  홀수일 때는 left가 right보다 1개 더 많아야 한다.
  (left에 중간값을 넣을 거니까 left의 길이가 1 큼)

 

left의 길이가 right의 길이+1 보다 길어진 경우

• 그런데,
  left의 최대값보다 같거나 작은 수는 left에 넣어버렸으니,
  값을 추가해가다보면 left가 right의 길이 차이가 1을 초과하는 경우가 생길 수 있다.

    # left가 right에 1을 더한 것보다 더 길면
    if len(left) > len(right) + 1:

• 이럴 때는 left의 가장 큰 값을 right로 옮겨줘야 한다.
  (left는 최대힙이므로, heappop을 했을 때 가장 큰 값이 나오고,
  right도 힙이니까 값을 넣으면 알아서 정렬이 되어 들어간다.)

        # left에서 제일 큰 값을 빼서 right에 넣습니다.
        temp = -heapq.heappop(left)
        heapq.heappush(right, temp)

 

right의 길이가 left의 길이보다 길어진 경우

• 반대로 right는 left보다 항상 길이가 같거나 작아야 하는데
  right가 더 길어질 경우도 생기겠지!

    # right가 left보다 길면
    elif len(right) > len(left):

• 그러면 right에서 제일 작은 값을 left에 넣어준다.

        # right에서 제일 작은 값을 빼서 left에 넣습니다.
        temp = heapq.heappop(right)
        heapq.heappush(left, -temp)

 

3. 출력

• 중간값이 left에 있으니까 left에서 가장 큰 값을 출력하면 된다!

    # 중간값은 최대힙에서 제일 큰 값입니다.
    print(-left[0])