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

문제

---

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

1655번: 가운데를 말해요

 

풀이

---

이 문제를 보고 가장 먼저 떠오른 방식은 입력을 받을 때마다 배열에 추가하고 배열을 정렬하여 가운데 값에 해당하는 인덱스를 출력으로 내보내는 것이었습니다. 하지만 문제에서 주어진 N의 최댓값과 시간 제한을 고려해보았을 때 위와 같은 방식은 시간초과가 날 것이 뻔하여 다른 풀이로 호다닥 전환했습니다.

 

문제에서 요구하는 가운데 숫자는 그 값을 기준으로 좌측과 우측 숫자의 개수가 동일하거나 좌측이 하나 적어야합니다.

(위에서 언급한 좌측과 우측은 각각 가운데 숫자를 기준으로 작은값들의 모임, 큰 값들의 모임을 말합니다. 또한 좌측이 하나 적은 경우는 숫자들의 총 개수가 짝수인 경우입니다.)

 

위와 같은 생각을 바탕으로 좌측은 최대힙, 우측은 최소힙을 사용하면서 좌,우 측의 크기를 잘 조정해주면 문제에서 원하는 답을 구할 수 있습니다. 

출처: https://art-coding3.tistory.com/44

 

 

백준이가 숫자를 외칠때마다 좌, 우측 중 적절한 힙에 숫자를 삽입하고 가운데 값을 조정합니다.

(+ 저의 풀이에서는 좌측 최대힙의 첫 번째 원소가 중간값이 되도록 코드를 썼습니다.)

• 우선 숫자 삽입 시 가장 중요하게 고려해야할 두 가지는 다음과 같습니다.
  
  • 삽입할 숫자와 가운데 값과의 대소 관계
  • 좌, 우측 힙의 밸런스

이 중 두 번째인 좌, 우측 힙의 밸런스에 대해 먼저 설명하면 아래와 같습니다.

• 좌측 최대힙에 중간값까지 들어간다는 사실을 고려하면 우측 힙은 좌측힙보다 크기가 커지면 안됩니다.
• 또한 좌측 힙 역시 우측 힙보다 2개 이상 크기가 커지면 안됩니다. 
• 위의 두 가지 조건을 위배한다면 좌측 힙의 첫번째 원소가 가운데 값이 아니게 됩니다.

 

숫자 삽입 시 첫번째 고려사항이었던 대소 관계에 따른 로직은 아래와 같습니다.

 

1. 백준이가 외친 숫자가 중간값보다 큰 경우

• 기본적으로 중간값보다 크므로 우측 최소힙에 해당 숫자가 들어가야합니다.
• 삽입 후에도 좌, 우측 힙의 밸런스가 무너지지 않는다면 가운데 값은 유지됩니다.
• 만약 삽입 후에 우측힙의 크기가 좌측힙보다 커진다면 즉, 좌, 우측 힙의 밸런스가 무너진다면 아래의 로직을 따라줍니다.
  • 우선 우측 힙에 숫자 삽입을 한다.
  • 우측 힙의 첫번째 원소를 뽑아서 좌측 힙에 추가해준다. ( 중간값보다 큰 값 중 최소 값을 뽑아 좌측으로 옮긴다. )
  • 이러면 좌측 힙의 첫번째 원소가 가운데 값이 됩니다!

2. 백준이가 외친 숫자가 중간값보다 작은 경우

• 중간값보다 작으므로 좌측 최대힙에 해당 숫자가 들어가야합니다.
• 숫자를 삽입하여도 양 힙 간의 밸런스가 무너지지 않는다면 좌측 힙에 숫자 삽입!
• 삽입을 후에 좌, 우측 힙의 밸런스가 무너질 것 같다면 숫자 삽입 전에 좌측 힙의 첫 번째 원소를  우측 힙으로 옮깁니다.
• 그 후 백준이가 외친 숫자를 좌측 힙에 삽입해줍니다.

위와 같은 과정을 반복하면 계속하여 좌측 힙의 첫 번째 원소가 문제에서 요구하는 가운데 값이 되도록 유지할 수 있습니다.

 

풀이 코드는 아래와 같습니다.

from heapq import heappush, heappop
import sys

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

leftPart, rightPart = [], []
leftCount, rightCount = 0, 0
midValue = 0

# 맨 처음 첫 값은 leftPart에 삽입
curNumber = int(input())
midValue = curNumber
# 최대힙 유지를 위해 음수값으로 변환해서 삽입
leftPart.append(-curNumber)
leftCount += 1
print(midValue)

for _ in range(n - 1):
    curNumber = int(input())
    # rightPart에 삽입해야하는 경우
    if curNumber >= midValue:
        # 양쪽 밸련스가 무너져 조정해야하는 경우
        if rightCount == leftCount:
            heappush(rightPart, curNumber)
            rightNumber = heappop(rightPart)
            heappush(leftPart, -rightNumber)
            leftCount += 1
        else:
            heappush(rightPart, curNumber)
            rightCount += 1

    elif curNumber < midValue:
        # 양쪽 밸련스가 무너져 조정해야하는 경우
        if leftCount - rightCount >= 1:
            leftNumber = heappop(leftPart)
            heappush(leftPart, -curNumber)
            heappush(rightPart, -leftNumber)
            rightCount += 1
        else:
            heappush(leftPart, -curNumber)
            leftCount += 1
    midValue = -leftPart[0]
    print(midValue)