[프로그래머스/Python] 디펜스 게임

디펜스 게임드디어 heap을 사용해 볼 기회가 찾아왔다..

문제접근

가장 먼저 생각해 볼 수 있는 방법은, 배열을 돌면서 무적권을 쓸지 정면 돌파를 할지 경우의 수를 나누는 방법이다.
이를 위해서 고민할 것은 시간복잡도..

• 1 ≤ n ≤ 1,000,000,000
• 1 ≤ k ≤ 500,000
• 1 ≤ enemy의 길이 ≤ 1,000,000
• 1 ≤ enemy[i] ≤ 1,000,000
  

보다시피 enemy의 길이가 너무 크다..2 가지의 경우를 나눠서 DFS 해주기엔 시간복잡도가 2^1,000,000이 나오게 된다..
그렇다면 최소 Nlog(N)의 복잡도를 사용해야 한다.

반복문을 한 번만 돌아한다는 말인데..여기서 고민 끝에 내린 결론이 heap을 사용하자였다.

1. enemy 반복문을 돈다.
2. enemy[i]가 n보다 작으면 맞서 싸운다. (answer+=1)
3. enemy[i]가 n보다 크면 무적권을 쓴다.
   (1) 이 전에 맞서 싸운 enemy중에, 더 큰 enemy가 있다면 그때 enemy에게 무적권을 쓴걸로 치자 ! (이를 위해 heappop을 쓴다)
   (2) 더 큰 enemy가 없다면, 지금 쓴다.
   (3) 첫 싸움이면 무조건 무적권을 쓴다.(무적권 없으면 끝 answer 출력)

내가 쓴 풀이

이해를 돕고자 리팩토링을 안 했습니다..위에 과정 그대로 해설 시작
참고로 여기서는 MaxHeap을 쓰고자, (-)를 곱해서 heappush를 했다.

import heapq
def solution(n, k, enemy):
    answer = 0
    h=[] # heap배열
    for i in enemy: # enemy 반복문을 돈다.
        if n>=i: # n보다 작으면 맞서 싸움
        
            n-=i
            heapq.heappush(h,-i) # 기록한다.
            answer+=1
            
        elif n < i and k != 0: # n보다 크면 무적권을 씀
        
            if h and -h[0]>=i: # 이 전에 맞서 싸운 더 큰 enemy가 있다면
                a=-heapq.heappop(h)
                n+=a # 그때 무적권을 썼던 것으로 친다.
                k-=1
                n-=i
                heapq.heappush(h,-i)
                answer+=1
            else: # 이 전에 맞서 싸운 더 큰 enemy가 없다면
                k-=1 # 지금 무적권을 쓴다.
                answer+=1
                
        elif n < i and k==0: return answer #무적권 없으면 끝 
    return answer

검색 풀이

from heapq import heappop, heappush

def solution(n, k, enemy):
    answer, sumEnemy = 0, 0
    heap = []

    for e in enemy:
        heappush(heap, -e)  # 현재 적을 힙에 추가
        sumEnemy += e       # 현재까지의 적의 합을 더함
        if sumEnemy > n:    # 적의 합이 n을 초과하면
            if k == 0:      # k가 0이면 더 이상 처리 불가, 루프 종료
                break
            sumEnemy += heappop(heap)  # 가장 큰 적을 되살림으로써 합을 줄임
            k -= 1         # k를 하나 소모
        answer += 1        # 적을 처리했으므로 answer를 증가
    return answer          # 최종 처리한 적의 수 반환

처리 시간은 비슷한데 훨씬 깔끔하다. 리팩토링 된 느낌인데, 실전에서 이렇게 깔끔한 코드를 생각하는 게 관건인 것 같다..
그래도 저번에 호텔 대실 문제를 풀면서 heap은 생각도 못 했는데, 이번 알고리즘 문제를 풀면서 heap을 어떤 식으로 활용해야 하는지 감이 온 것 같다.