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

off_sujin·2022년 6월 9일
0

BOJ 단기간 성장

목록 보기
2/2

문제

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

문제 풀이

이 문제의 의도는 가운데 값을 찾기 위해 우선순위 큐를 사용하는 것이다.

최대힙과 최소힙, 이렇게 우선순위 큐를 2개 사용한다.

최대힙에서 가장 큰 원소와 최소힙에서 가장 작은 원소를 비교하여 가운데 수를 구할 수 있다.

우선순위 큐를 구현하려면 파이썬의 heapq를 사용하면 된다.

그런데 이 모듈은 최소힙이기 때문에 각 원소를 음수로 바꿔서 최대힙으로 만들어주면 된다.

알고리즘

가운데 값은 항상 최대힙의 첫번째 원소라고 가정한다.

  1. N번 반복문을 돌며 숫자를 입력받는다.

  2. 지금까지 힙에 저장된 숫자의 개수가 짝수면 숫자를 최대힙에 넣고, 홀수면 최소힙에 넣는다.

  3. 최소힙에 원소가 있고, 최소힙의 최댓값보다 최대힙의 최솟값이 작은 경우 둘을 바꿔준다.

  4. 중간값인 최대힙의 첫번째 원소를 출력한다.

최소힙에만 원소가 있는 경우(n=1인 경우)가 존재하기 때문에, 3번에 최대힙에 원소가 존재하는지 확인하는 조건문을 추가했다.

입력으로 [1, 5, 10, 2, -99, 7] 가 들어온다면

n = 1, input = 1
지금까지 힙에 저장된 숫자의 개수는 0개로 짝수이기 때문에 왼쪽에 있는 최대힙에 넣는다.
최소힙에 원소가 없기 때문에 중간값은 최대힙의 0번째 원소인 1이다.

n = 2, input = 5
지금까지 힙에 저장된 숫자의 개수는 1개로 홀수이기 때문에 오른쪽에 있는 최소힙에 넣는다.
최소힙에 원소가 있으므로 최대힙의 최댓값인 1과 최소힙의 최솟값인 5를 비교한다.
최솟값이 최댓값보다 크므로 그냥 내비둔다.
이번에도 중간값은 최대힙의 0번째 원소인 1이다.

n = 3, input = 10
지금까지 힙에 저장된 숫자의 개수는 2개로 짝수이기 때문에 왼쪽에 있는 최대힙에 넣는다.
최소힙에 원소가 있으므로 최대힙의 최댓값인 10과 최소힙의 최솟값인 5를 비교한다.
최댓값이 최솟값보다 크므로 둘을 바꿔준다.
이제 중간값은 최대힙의 0번째 원소인 5이다.

위에 적인 숫자는 힙의 인덱스이다!

n이 6이 될 때까지 위의 과정을 반복한다.

n = 4, input = 2

중간값 = 2

n = 5, input = -99

중간값 = 2

n = 6, input = 7

중간값 = 2
n이 6이 되었으므로 반복문을 빠져나온다.

소스 코드

import heapq
import sys

input = sys.stdin.readline

N = int(input())
left_h = []  # 최대힙, 첫번째 원소는 항상 중간값이다.
right_h = []  # 최소힙

for i in range(N):
    num = int(input())

    # 숫자의 개수가 짝수면 왼쪽에 넣고 홀수이면 오른쪽에 넣는다.
    if i % 2 == 0:
        heapq.heappush(left_h, -num)
    else:
        heapq.heappush(right_h, num)

    # 오른쪽에 원소가 있고, 왼쪽의 최댓값보다 오른쪽의 최솟값이 작은 경우 둘을 바꿔준다.
    if right_h and (-left_h[0] > right_h[0]):
        min_value = heapq.heappop(right_h)
        max_value = -heapq.heappop(left_h)

        heapq.heappush(right_h, max_value)
        heapq.heappush(left_h, -min_value)

    # 중간값을 출력한다.
    print(-left_h[0])

입력에 대한 최소힙과 최대힙을 출력한 결과는 다음과 같다.

profile
학습 중..

0개의 댓글