[Python/Heap] 디펜스 게임

문제 설명

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 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	key	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라운드까지 막을 수 있습니다.

---

풀이

1.
물리칠 수 있는 판 중 적의 수가 큰 판에서 '무적권'을 사용해야 하므로
최솟값을 쉽게 pop할 수 있는 heap을 사용해야 한다.

최댓값을 pop해야 하므로 -를 붙여서 push를 하고, 적을 싸우면서 n,k 값에서 빼는 게 아니라 tmp 값에 물리친 적의 수를 더한다.
그래서 tmp가 n보다 커질 때 무족권을 써야 하는 상황이고 쓸 수 있다면 tmp에서 최댓값을 빼주고, 무족권을 쓸 수 없다면 값을 반환한다.

'''
병사 n명. 매 라운드마다 적이 enemy[i]명
남은 병사 수가 현재 라운드의 적 수보다 적으면 게임 종료
무적권 사용하면 병사 사용 않고 라운드 이김(k번 사용 가능)

몇 라운드까지 막을 수 있는 지 return
enemy에서 큰 값을 차례로 무적권 사용 -> 그 때까지 못 가면?
'''
from heapq import heappush, heappop
def solution(n, k, enemy):
    answer = 0
    heap = [] # 힙에는 무족권을 쓸 경우를 담는다
    tmp = 0 # 싸운 적의 수
    
    for e in enemy:
        heappush(heap, -e) # 힙은 최솟값부터 pop하므로 -
        tmp += e
        if tmp > n: # 무족권을 써야 함
            if k == 0: # 무족권 다 씀
                break
            else: # 무족권 사용 가능
                tmp += heappop(heap) #음수값이 나오므로
                k -= 1
        #print(heap, tmp)
        answer += 1
                
    return answer

from heapq import heappush, heappop
def solution(n, k, enemy):
    answer = 0
    heap = [] # 힙에는 무족권을 쓸 경우를 담는다
    tmp = 0 # 싸운 적의 수
    
    for e in enemy:
        heappush(heap, -e) # 힙은 최솟값부터 pop하므로 -
        tmp += e
        if tmp <= n: # 무족권 안 써도 됨
            answer += 1
        elif k > 0: # 무족권 써야 하고 쓸 수 있음
            k -= 1
            tmp += heappop(heap)
            answer += 1
        print(heap, tmp)            
    return answer
'''
첫번째 테케
[-4] 4
[-4, -2] 6
[-4, -2] 6
[-4, -2] 6
[-3, -2] 5
[-3, -2, -3] 8
[-3, -2, -3, -1] 9
 '''

 

2.
무족권이 k개라는 의미는 k번째 라운드까지는 이기는 것을 보장한다는 의미이다.

# k번째까지는 무조건 무찌른다

import heapq 
def solution(n, k, enemy):
    q = enemy[:k] 
    heapq.heapify(q)
    for idx in range(k,len(enemy)):
        n -= heapq.heappushpop(q,enemy[idx]) 
		# enemy[idx] 삽입하고 이를 반영한 데서 최솟값을 n에서 뺌
        if n < 0:
            return idx
    return len(enemy)

 

📝 heappushpop
push를 한 후에 이를 반영한 데서 최솟값을 pop함

from heapq import heappush, heappushpop

heap = []
heappush(heap, 1)
heappush(heap, 2)
heappush(heap, 3)
print(heap) # [1, 2, 3]
n=heappushpop(heap, 4)
print(heap) # [2, 4, 3]
print(n) #1

from heapq import heappush, heappushpop

heap = []
heappush(heap, 1)
heappush(heap, 2)
heappush(heap, 3)
print(heap) # [1, 2, 3]
n=heappushpop(heap, 0)
print(heap) # [1,2,3]
print(n) #0

 

 

틀린 풀이

# *** 오답: 언제 k를 쓸 지 몰라서 이전 판의 병사보다 큰 경우라고 임시 지정 ***
def solution(n, k, enemy):
    answer = -1
    tmp = -1
    for e in enemy:
        if k == 0 and n <= 0:
            return answer
        if k > 0 and n <= 0: # n을 다써서 무적권를 사용하는 경우
            k -= 1
        elif e >= tmp: # 병사 수가 커서 무적권를 사용하는 경우
            tmp = e
            k -= 1
        else: # 직접 싸우는 경우
            n -= e
        answer += 1
        #print(n, k)
    return answer

# *** 오답 ***
# 일단 직접 다 싸우고 무족권 씀
from heapq import heappush, heappop
def solution(n, k, enemy):
    answer = 0
    for e in enemy:
        if k == 0 and n <= 0:
            break
        if n >= e: # 직접 싸우는 경우
            n -= e
        else:
            if k == 0: # 무족권 다 쓴 경우
                break 
            else: # 무족권 쓸 차례
                k -= 1
        answer += 1
        #print(n, k)
    return answer