• TOC {:toc}

이 글은 백준 온라인 저지의 1655번 문제를 파이썬(Python)으로 풀이한 것을 모아놓은 글입니다.

일종의 연습 기록이며 제가 정답을 받은 코드와 참고할만한 다른 코드를 같이 기록합니다. 필요한 경우 코드에 대한 해설을 기록합니다만 코드는 통과했어도 해설은 틀릴 수 있기 때문에 가볍게 참고해주시길 부탁드립니다. 피드백은 편하신 방법으로 자유롭게 주시면 감사하겠습니다.

2021.04.29

메모리 시간 코드 길이
KB ms B
단계 시작 시각 끝난 시각 걸린 시간
문제 이해 20:55:36 20:56:10
풀이 생각 1 20:56:11 20:58:25
코딩 1 20:58:37 21:00:51
풀이 생각 2 21:02:28 22:40:15
코딩 2 22:40:17 23:29:29
import sys
from heapq import heappush, heappop
input = lambda: sys.stdin.readline().rstrip()

S = [] # 최대 힙
L = [] # 최소 힙
tmp = []

for i in range(int(input())):
    # 입력을 받아 tmp에 넣는다.
    tmp.append(int(input()))
    # S에 원소가 존재하면
    if S:
        # pop 해서 tmp에 넣는다.
        tmp.append(-heappop(S))
    # L에 원소가 존재하면
    if L:
        # pop 해서 tmp에 넣는다.
        tmp.append(heappop(L))
    # tmp를 정렬한다.
    tmp.sort()

    # tmp의 최솟값을 S에 넣는다.
    heappush(S, -tmp.pop(0))
    # tmp에 원소가 남아있으면
    if tmp:
      # tmp의 최댓값을 L에 넣는다.
        heappush(L, tmp.pop())
    # tmp에 원소가 남아있고
    if tmp:
        # S의 원소가 L보다 많으면
        if len(S) > len(L):
            # tmp에 남아있는 값(중간값)을 L에 넣는다.
            heappush(L, tmp.pop())
        # S의 원소의 수가 L과 같으면
        else:
            # tmp에 남아있는 값을 S에 넣는다.
            heappush(S, -tmp.pop())
    
    print(-S[0])

아이디어 & 풀이

최소 힙과 최대 힙을 둘 다 이용해야 한다.

  • 중간값보다 작은 값을 모아두는 최대 힙 S: 중간값이 최댓값
  • 중간값보다 큰 값을 모아두는 최소 힙 L: 중간값이 최솟값

수를 입력받으면 S의 최댓값과 L의 최솟값을 받아온 뒤 세 수를 비교해서 다시 SL에 재분배한다.

  • S의 원소 개수와 L의 원소 개수가 같은 경우엔 S로 그렇지 않으면 L로 push 한다.
  • 이 경우 마지막에 S에서 pop 한 최댓값이 중간값이다.

예제 1의 진행 과정은 다음과 같다.

  • 입력: [1, 5, 2, 10, -99, 7, 5]
입력 pop S pop L 정렬 S L
1 1 [1] []
5 [], 1 [] 1, 5 [1] [5]
2 [], 1 [], 5 1, 2, 5 [1, 2] [5]
10 [1], 2 [], 5 2, 5, 10 [1, 2] [10, 5]
-99 [1], 2 [10], 5 -99, 2, 5 [-99, 1, 2] [10, 5]
7 [-99, 1], 2 [10], 5 2, 5, 7 [-99, 1, 2] [10, 7, 5]
5 [-99, 1], 2 [10], 5 2, 5, 5 [-99, 1, 2, 5] [10, 7, 5]

디버그

  • 힙으로 (H) 받은 다음에 힙의 중간인 H[len(H) - 1 // 2]를 출력하려고 했는데 힙이 순차 정렬된 리스트가 아니었다.
  • 결국 검색해서 찾아봤다. 최소 힙과 최대 힙을 둘 다 이용해야 한다.

피드백

  • 매번 SL에서 pop 해서 비교할 필요 없이, 일단 순서에 맞게 SL에 집어넣고 두 값을 비교해 필요한 경우에만 pop 해 반대로 push 하면 됐었다 (풀이 1).
  • 입력받은 값을 대소 관계에 맞게 넣기 위해 SL 모두와 비교할 필요가 없었다 (풀이 2).
    • 원소를 추가해야 하는 힙이 무엇인가에 따라 두 힙 중 하나의 root 값과만 비교하면 된다.
    • 예를 들어 S에 원소를 추가로 push 해야 할 때는 L의 root와 입력값 v만 비교하면 된다.
  • SL에 번갈아 가면서 값을 추가해야 하기 때문에 SL의 길이를 비교하는 것 대신 반복문에서 i를 받아와서 i 값에 따라 경우를 나눠도 된다.

참고 답안 1

# 풀이 1-1
import sys
from heapq import heappush, heappop
input = sys.stdin.readline

S, L = [], [] # S: 최대 힙, L: 최소 힙
for i in range(int(input())):
    # i가 짝수이면
    if i % 2 == 0:
        # S에 값을 집어 넣는다
        heappush(S, -int(input()))
    # 홀수이면
    else:
        # L에 값을 집어 넣는다.
        heappush(L, int(input()))

    # S의 root(최댓값)가 L의 root(최솟값)보다 크면
    if L and -S[0] > L[0]:
        # 두 값을 바꾼다.
        heappush(S, -heappop(L))
        heappush(L, -heappop(S))

    print(-S[0])

# 풀이 1-2
# 풀이 1-1을 축약 표현과 heapreplace를 이용해서 더 짧게 구현했다.
import sys
from heapq import heappush, heapreplace
input()

S, L = [], []
for i, x in enumerate(sys.stdin):
    heappush(*([L, int(x)] if i % 2 else [S, -int(x)]))
    if l and r and -S[0] > L[0]:
        heapreplace(L, -heapreplace(S, -L[0]))
    print(-S[0])

아이디어 & 풀이

readline()을 사용하지 않고 for문에 sys.stdin을 돌려 stdin에 입력이 있을 때마다 실행하게 할 수 있다.

  • 처음 입력받는 N은 쓸 일이 없기 때문에 input()만 작성하고 끝낸다.

참고 답안 2

# 풀이 2-1
import sys
from heapq import heappush, heapreplace
input = sys.stdin.readline

S, L = [], []
for _ in range(int(input())):
    x = int(input())
    # S와 L이 원소 수가 같고
    if len(S) == len(L):
        # L(S)에 원소가 없거나
        # L의 root(최솟값)가 v보다 크거나 같으면
        if len(L) == 0 or L[0] >= x:
            # S에 v를 집어넣는다.
            heappush(S, -x)
        # L의 root가 v보다 작으면
        else:
            # S에 L의 root를 넣고
            heappush(S, -L[0])
            # L의 root를 v로 바꾼다.
            heapreplace(L, x)
    # S의 원소 수가 L보다 많고
    else:
        # S의 root(최댓값)이 v보다 작거나 같으면
        if -S[0] <= x:
            # L에 v를 집어 넣는다.
            heappush(L, x)
        # S의 root가 v보다 크면
        else:
            # L에 S의 root를 넣고
            heappush(L, -S[0])
            # S의 root를 v로 바꾼다.
            heapreplace(S, -x)
    # S의 root를 출력한다.
    print(-S[0])

# 풀이 2-2
import sys
from heapq import heappop, heappush
input = sys.stdin.readline
L = []
S = []
med = 0

for i in range(int(input())):
    x = int(input())
    # i가 0일 때는
    if i == 0:
        # 입력값이 중간값이다.
        med = x
    # i가 짝수일 때는
    # (S의 원소 개수가 증가할 차례)
    elif i % 2:
        # 입력받은 값이 중간값보다 크면
        if x > med:
            # L에 입력받은 값을 집어넣고
            heappush(L, x)
            # S에 중간값을 집어넣고
            heappush(S, -med)
            # 중간값을 L의 최솟값을 pop 한 값으로 바꾼다.
            med = heappop(L)
        # 입력받은 값이 중간값보다 작으면
        else:
            # S에 입력받은 값을 집어넣는다.
            heappush(S, -x)
    # i가 홀수일 때는
    # (L의 원소 개수가 증가할 차례)
    else: 
        # 입력받은 값이 중간값보다 작으면
        if x < med:
            # S에 x를 집어넣고
            heappush(S, -x)
            # 현재 중간값을 L에 집어넣고
            heappush(L, med)
            # 중간값을 S의 최댓값을 pop 한 값으로 바꾼다.
            med = -heappop(S)
        # 입력받은 값이 중간값보다 크면
        else:
            # L에 입력받은 값을 집어넣는다.
            heappush(L, x)
    # 중간값을 출력한다.
    print(med)

# 풀이 2-3
from sys import stdin
from heapq import heappush, heappop

input()
S, L = [], []
for x in map(int, stdin):
    # S가 존재하지 않거나
    # 입력받은 값이 S의 root(최댓값)보다 작거나 같으면
    if not S or x <= -S[0]:
        # 입력값을 S에 집어넣고
        heappush(S, -x)
        # S의 원소가 L의 원소보다 두 개 이상 많으면
        if len(S) > len(L) + 1:
            # S의 root를 pop 해서 L에 push 한다.
            heappush(L, -heappop(S))
    # 입력받은 값이 S의 root보다 크면
    else:
        # 입력받은 값을 L에 집어넣고
        heappush(L, x)
        # L의 원소가 S의 원소보다 많으면
        if len(S) < len(L):
            # L의 root를 pop 해서 S에 push 한다.
            heappush(S, -heappop(L))

    print(-S[0])

아이디어 & 풀이

각 풀이가 다른 방법으로 보이긴 하지만 모두 다음 두 가지로 경우를 나누고 있다.

  • SL의 원소 개수: SL의 원소 개수는 같거나 S가 ‘1개’ 더 많다.
  • 입력값과 S, L root 값의 대소관계: S의 root(최댓값) <= 중간값 < L의 root(최솟값)을 유지해야 한다.

각 경우에 따른 push / pop을 정리하면 다음과 같다.

  • medS에 추가할 때는 L[0], L에 추가할 때는 -S[0]로 잡아주면 된다.

  • SL에 pop과 push를 둘 다 하는 경우 순서에 상관없이 pop 되는 원소는 동일하기 때문에 순서는 크게 상관없다.

    • 예를 들어 왼쪽 아래의 경우 x < -S[0]xS의 root(최댓값)보다 작기 때문에 먼저 push 해도 이후에 pop 될 root가 바뀌지 않는다.
    x < med x > med
    len(S) == len(L): S에 원소를 추가 heappush(S, -x) heappush(S, -heappop(L)), heappush(L, x)
    len(S) > len(L) + 1: L에 원소를 추가 heappush(L, -heappop(S)), heappush(S, -x) heappush(L, x)
  • S에 원소가 없을 경우를 따로 생각해줘야 한다.

    • if not S if len(S) == 0
    • 이 경우 조건을 S에 입력값을 push 해주는 경우(i.e., heappush(S, -x))에 or로 같이 묶어주면 된다.

S의 root 값을 출력하는 대신 중간값 med 변수를 만들어 관리해도 된다.