[프로그래머스] 무지의 먹방 라이브 - Python

프로그래머스

* 효율성 테스트에 부분 점수가 있는 문제입니다. 라고 나와있으니 문제 그대로 풀면 분명 시간 초과가 나올겁니다.

POINT

시간이 적게 걸리는 음식부터 공략한다!

• k를 초과하지 않는 선에서, 먹는 시간이 적은 음식부터 먹어치운다. → 시간이 적은 음식부터 나열하기 위해 우선순위 큐를 사용해보자.
• 한 음식을 다 먹는 시간이 k를 초과한다면, 원래 풀이대로 먹고 남은 시간만큼 번호를 센 후 반환한다.

시간은 작은 순으로 2번 → 3번 → 1번을 먹는다고 하면, 해당 음식을 다 먹는데 걸리는 시간은 위와 같습니다. 주의할 점은 이전 음식을 먹는데 걸린 시간을 빼줘야 시간이 중복되지 않는다는 것입니다.

우선순위 큐 heapq

heapq.heappush(heap, item)

• heap에 item 원소를 넣는다.

heapq.heappop(heap)

• heap에서 우선순위가 높은 원소를 가져옵니다.

풀이

import heapq

def solution(food_times, k):
    # 음식을 다 먹는 경우
    if sum(food_times) <= k:
        return -1
    
    q = []
    # (먹는 시간, 음식 번호)를 우선순위 큐에 저장
    for i in range(len(food_times)):
        heapq.heappush(q, (food_times[i], i+1))
    
    length = len(food_times) # 남은 음식 개수
    previous = 0 # 이전 음식을 먹은 시간
    
    # 우선순위가 가장 높은 음식을 다 먹을 수 있으면
    while (q[0][0] - previous) * length <= k:
        now = heapq.heappop(q)[0] # 현재 음식 시간
        k -= (now - previous) * length  # 남은 시간
        length -= 1 # 다 먹은 음식 개수 빼주기
        previous = now # previous += now - previous
        
    # 먹고 남는 시간은 직접 센다.
    result = sorted(q, key=lambda x: x[1]) # 음식 번호 순으로 정렬
    return result[k % length][1] # 다음 음식 번호 반환

 

참고자료

heapq — 힙 큐 알고리즘 — Python 3.10.5 문서

 

이것이 취업을 위한 코딩 테스트다 with 파이썬 - YES24