[프로그래머스] 힙(Heap) - 라면공장 (python)

힙(Heap) - 라면공장, 문제확인

개인적으로 우선순위 큐를 어떻게 활용해야 하는지 전혀 감이 잡히지 않아 헤맸던 문제이다. 일단, 문제 설명은 아래와 같다.

 

---

라면 공장은 하루에 밀가루 1톤 사용
k일 이후에 밀가루 공급 가능, 해외 공장에서 수입해야 함
해외 공장에서는 밀가루 공급이 가능한 날짜와 수량을 알려줌
라면 공장에서는 운송비를 줄이기 위해 최소한의 횟수로 밀가루를 공급받고자 함

현재 공장에 남아있는 밀가루 수량 : stock 
밀가루 공급 일정 :dates (dates[i] " i번째 공급 가능일)
해당 시점에 공급 가능한 밀가루 수량 : supplies (supplies[i] : dates[i] 날짜에 공급 가능한 밀가루 수량)
원래 공장으로부터 공급받을 수 있는 시점 : k

밀가루가 떨어지지 않고 공장을 운영하기 위해 최소한 몇번 밀가루를 공급받아야 하는지 리턴

밀가루가 바닥나는 경우는 주어지지 않음

 

 

제시되어 있는 예를 보면, stock=4, dates=[4,10,15], supplies=[20,5,10], k=30일 때, 결과 값은 2가 된다. 현재 밀가루가 4톤 남아있기 때문에 4일째에 반드시 20을 공급받게된다. 11일 후인 15일째에는 9톤의 밀가루가 남아있게 되고 이때 10톤을 더 공급받으면 19톤이 남아있게 되고 30일째까지 버틸 수 있으므로 총 2회만 공급받으면 되는 것이다.

 

 

힙(Heap) - 라면공장, 문제풀이

이게 설명 이해는 간단한데 막상 코드로 작성하려니 여간 복잡한게 아니었다. 게다가 dates와 supplies는 서로 인덱스가 일치해야 하기 때문에 우선순위 큐를 어떻게 적용시켜야할지도 난감했다. 중요한 포인트는 버틸 수 있는 날짜 안에서 제일 큰 공급물량을 선택하는 것이다.

 

먼저, heapq는 최소힙 기반으로 구현되어있는 모듈이기 때문에 최대힙을 써야하는데 다른 사람 코드를 참조해보니 아예 모든 원소에 -를 붙여 저장을 하는 방식을 사용하길래 그 방식을 차용했다. 이렇게되면 최솟값을 꺼낸것에 -1을 곱해준 값이 최대값이 되기 때문이다.

 

반복문을 총 2개를 도는데 첫번째 반복문은 쌓아둔 재고가 k를 넘을 때까지 반복한다. 여기서 코드 중간에 재고를 갱신하면서 계산을 할텐데 예를 들어, 4일째에 재고가 0이되고 20을 추가로 공급받을 때 갱신된 재고값은 20이 아니라 24로 계산했다. 남아있는 날 수가 아닌 절대적 날 수를 세야 k랑 비교하기 편하기 때문이다. 

 

import heapq

def solution(stock, dates, supplies, k):
    q=[]
    count=0
    while stock<k: # k날 까지 버티기 충분한 재고가 쌓이면 멈춤
        while 1:
            if len(dates)==0 or dates[0]>stock: break #dates 길이가 0이거나 dates[0]이 stock보다 크면 멈춤
            d=dates.pop(0) #dates[0]을 꺼냄
            s=supplies.pop(0) #supplies[0]을 꺼냄
            heapq.heappush(q, -s) #우선순위 큐에 저장     
        ms=(heapq.heappop(q) * -1) #우선순위큐에서 최대값을 꺼냄
        stock=stock+ms #stock에 최대값을 더해 stock값을 갱신해줌
        count+=1 #count를 +1해줌
    return count

 

 

하지만 한 테스트케이스에서 시간초과가 발생했다.

 

 

 

이중 포문은 못고칠 것 같고, 다른 사람의 풀이를 참조하니 매번 조건문에서 dates의 길이를 체크하고 pop을 해주는 것이 시간 초과가 생긴 원인인 것 같았다. 코드가 조금 더럽긴 하지만, i 인덱스를 사용해서 pop 부분만 제거해주니 시간초과를 해결할 수 있게됐다.

 

import heapq

def solution(stock, dates, supplies, k):
    q=[]
    count=0
    i=0
    while stock<k: 
        while 1:
            if len(dates)==i or dates[i]>stock: break
            d=dates[i]
            s=supplies[i]
            i+=1
            heapq.heappush(q, -s)     
        ms=(heapq.heappop(q) * -1)
        stock=stock+ms
        count+=1
    return count

 

 

으 힙은 확실히 이전 스택/큐나 해시에 비해 어려운 느낌이 든다.