[백준/Python]2461: 대표 선수

문제

https://www.acmicpc.net/problem/2461
백준 문제집 - 0x14강 - 투 포인터

 

알고리즘: 투 포인터, 우선순위 큐, 정렬

---

풀이

투포인터만 사용하면 시간초과가 나는 문제다.
Python3로 통과하기 위해선 우선순위 큐, 힙heapq을 사용해야한다.

 

우선 각 반의 능력치를 모두 오름차순 정렬한다.
그리고 힙 리스트 heap에 각 반의 최솟값들을 (값, 반, 인덱스) 형태로 삽입한다.
이 중 가장 큰 값은 따로 저장해둬야 한다.

1. heap에서 최솟값을 꺼냄 => (값, 반, 인덱스)
2. 기존 결과값과 (최댓값 - 최솟값) 중 더 작은 값을 새로운 결과값으로 업데이트.
3. 인덱스를 1 증가시킴.
   • 해당 반에서 몇번째 값을 가리키는지 나타내는 포인터 역할
   • 즉, 현재까지의 가장 최솟값의 포인터를 뜻함.
   • +1 처리 후 인덱스가 M이 된다면, 더 이상 탐색할 포인터가 없다는 뜻이므로 break.
4. 꺼낸 반의 다음 최솟값과 포인터를 삽입 => (다음 값, 반, 1 증가시킨 인덱스)
5. 기존 최댓값과 새로 집어넣은 최솟값을 비교 후 더 큰값을 새로운 최댓값으로 업데이트

 

전체 코드

# 메모리: 70676KB / 시간: 1448ms
from sys import stdin
from heapq import heappop, heappush


input = stdin.readline

N, M = map(int, input().split())
classes = [sorted(map(int, input().split())) for _ in range(N)]

heap = []
max_stat = 0

# 각 반의 최솟값을 heap에 추가
for i in range(N):
    heappush(heap, (classes[i][0], i, 0))
    max_stat = max(classes[i][0], max_stat)

min_diff = float("inf")

while True:
    # 현재 heap에서 최솟값을 꺼냄
    min_stat, class_idx, pos = heappop(heap)

    min_diff = min(max_stat - min_stat, min_diff)

    pos += 1  # 포인터를 1 증가시킴
    if pos == M:
        break

    # 해당 반의 다음 최솟값과 포인터를 삽입
    heappush(heap, (classes[class_idx][pos], class_idx, pos))
    max_stat = max(classes[class_idx][pos], max_stat)


print(min_diff)

 

참고로 최댓값을 미리 구해놓지 않으면 시간초과에 걸린다.
매 턴마다 heap의 모든 요소를 탐색해야하기 때문이다.

 

그리고 이번 문제를 풀며 힙 사용시 주의해야할 점을 알게되었다...
힙의 0번째 인덱스값은 항상 최솟값을 나타내지만, 1번째값이 그 다음 최솟값을 나타내는것은 아니다.
힙은 항상 완전이진트리 구조로 구성되어 있으므로 형제 노드간의 순서는 보장되지 않는다.

예를들어, [0, 5, 3, 9, 2]를 힙에 추가하게 되면 아래와 같다.

from heapq import heappush, heappop


arr = [0, 5, 3, 9, 2]
heap = []

for a in arr:
    heappush(heap, a)

print(heap)
# [0, 2, 3, 9, 5]

heappop(heap)  # 0 추출

print(heap)
# [2, 5, 3, 9]

이걸 제대로 이해하지 못한 채 최댓값을 heap[-1]로 구했었다가 틀림^^;
자료구조 공부를 다시 해야될듯싶다.

 

힙의 삽입/삭제 궁금하다면 아래 블로그를 참고하자.

https://heytech.tistory.com/69

[자료구조] 힙(Heap) 자료구조에 대해 알아보자!(+Python 구현)