[프로그래머스] 디펜스 게임 (Python) - heapq, heappush

문제 설명

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 n명으로 연속되는 적의 공격을 순서대로 막는 게임입니다. 디펜스 게임은 다음과 같은 규칙으로 진행됩니다.

• 준호는 처음에 병사 n명을 가지고 있습니다.
• 매 라운드마다 enemy[i]마리의 적이 등장합니다.
• 남은 병사 중 enemy[i]명 만큼 소모하여 enemy[i]마리의 적을 막을 수 있습니다.
  • 예를 들어 남은 병사가 7명이고, 적의 수가 2마리인 경우, 현재 라운드를 막으면 7 - 2 = 5명의 병사가 남습니다.
  • 남은 병사의 수보다 현재 라운드의 적의 수가 더 많으면 게임이 종료됩니다.
• 게임에는 무적권이라는 스킬이 있으며, 무적권을 사용하면 병사의 소모없이 한 라운드의 공격을 막을 수 있습니다.
• 무적권은 최대 k번 사용할 수 있습니다.

준호는 무적권을 적절한 시기에 사용하여 최대한 많은 라운드를 진행하고 싶습니다.

준호가 처음 가지고 있는 병사의 수 n, 사용 가능한 무적권의 횟수 k, 매 라운드마다 공격해오는 적의 수가 순서대로 담긴 정수 배열 enemy가 매개변수로 주어집니다. 준호가 몇 라운드까지 막을 수 있는지 return 하도록 solution 함수를 완성해주세요.

---

제한사항

• 1 ≤ n ≤ 1,000,000,000
• 1 ≤ k ≤ 500,000
• 1 ≤ enemy의 길이 ≤ 1,000,000
• 1 ≤ enemy[i] ≤ 1,000,000
• enemy[i]에는 i + 1 라운드에서 공격해오는 적의 수가 담겨있습니다.
• 모든 라운드를 막을 수 있는 경우에는 enemy[i]의 길이를 return 해주세요.

 

입출력 예

n	k	enemy	result
7	3	[4, 2, 4, 5, 3, 3, 1]	5
2	4	[3, 3, 3, 3]	4

---

입출력 예 설명

입출력 예#1

• 1, 3, 5 라운드의 공격을 무적권으로 막아내고, 2, 4 라운드에 각각 병사를 2명, 5명 소모하면 5라운드까지 공격을 막을 수 있습니다. 또, 1, 3, 4번째 공격을 무적권으로 막아내고, 2, 5 번째 공격에 각각 병사를 2명, 3명 소모하여 5라운드까지 공격을 막을 수 있습니다. 그보다 많은 라운드를 막는 방법은 없으므로 5를 return 합니다.

입출력 예#2

• 준호는 모든 공격에 무적권을 사용하여 4라운드까지 막을 수 있습니다.

---

from heapq import heappush, heappop

def solution(n, k, enemy):
    answer = 0
    enemy_len = len(enemy)
    if k >= enemy_len:
        return enemy_len
    
    idx = 0
    enemy_list = []
    while n >= 0 and idx < enemy_len:
        enemy_now = enemy[idx]
        if enemy_now <= n: #방어 가능 시 방어
            n -= enemy_now
            heappush(enemy_list, -enemy_now) # - 사용해 최대 힙 구현
        else: #방어 불가능 시, 
            if k > 0 and not enemy_list: #방어 기록 없을 시, 무적권 사용
                k -= 1      
            else: #방어했던 기록 있을 시, 막았던 것들 + 현재 값 중 가장 큰 값에 무적권 사용
                heappush(enemy_list, -enemy_now)
                enemy_list_max = -heappop(enemy_list) 
                if k > 0 and n + enemy_list_max - enemy_now >= 0:
                    n += enemy_list_max - enemy_now
                    k -= 1
                else:
                    break
        #업데이트
        idx += 1
        answer += 1
    
    return answer

힙을 쓰지 않으면 시간초과로 풀 수 없었던 문제 

= 힙을 사용하는 방법을 잘 모르는 내가 절대 풀 수 없었던 문제

 

1시간을 끙끙댔지만, 런타임에러도 실패도 아니고 시간초과가 뜨는 바람에 

머리를 붙잡고 ㅋㅋㅋㅋ 어떻게든 효율적으로 짜보려고 난리를 쳤지만... 결국 질문하기를 눌렀다

힙을 써야 한다는 말에 아 이건 내가 풀 수 없었던 문제구나 ㅎ 싶어서 바로 검색행 ...

힙을 만드는 방법만 공부하고 다시 풀었다

 

****중요했던 점:

최대 힙은 Python으로 만들어내는 옵션이 없으므로,

push할 때 -를 붙여서 음수로 만든 후 최소 힙을 구성해주어야 한다는 것.

그리고 pop 할 때 -를 붙여서 가져오면 원래 값을 얻을 수 있다!

 

자바로 구조 짤 땐 미칠듯이 어려웠던 힙이... 그저 heappush 하나로 해결이 된다

파이썬 라이브러리는 진짜 사기야

여전히 힙에 대한 것도 헷갈리는데 아주 손쉽게 힙을 만들었다.

 

그래서 정리해본다 다시 ....

자바로 공부할 때도 우선순위큐는 진짜 너무 어려웠는데... 이렇게 또 마주치네

잊지말자 힙

 

HEAP
- 완전 이진 트리
- 우선순위 큐 구현에 사용 (PriorityQueue)
- 최소값, 최대값 탐색에 효과적
Min. heap: 부모 노드의 키 값이 자식 노드의 키 값보다 작거나 같은 완전 이진 트리
Max. heap: 부모 노드의 키 값이 자식 노드의 키 값보다 크거나 같은 완전 이진 트리

*리스트로 직접 구현 시, 노드 인덱스
root node index = 1
left child node index = parent node index * 2
right child node index = parent node index * 2 + 1
parent node index = child node index // 2

*원소 삽입 시
1) 리스트의 맨 뒤에 append 후, 부모로 거슬러 올라가면서 swap
*최소 힙에서 최솟값, 최대 힙에서 최댓값 삭제 시
1) 리스트의 루트 노드와 마지막 원소를 swap 후 pop
2) 그 후 루트 노드와 자식을 비교하며 swap 해서 힙 유지

* 트리의 높이만큼만 비교를 진행하므로 최솟값, 최댓값을 찾는 것에 O(logN)의 복잡도를 가짐