[프로그래머스] Lv2. 디펜스 게임

디펜스 게임 

문제 설명

준호는 요즘 디펜스 게임에 푹 빠져 있습니다. 디펜스 게임은 준호가 보유한 병사 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 해주세요.

---

입출력 예nkenemyresult

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라운드까지 막을 수 있습니다.

틀린 풀이

def solution(n, k, enemy):
    answer = 0
    combs = []
    
    hap = 0
    k_cnt = 0
    for e in enemy:
        hap += e
        combs.append(e)
        
        if hap > n:
            if k_cnt < k:
                m = max(combs)
                hap -= m
                k_cnt += 1
                combs.remove(m)
            else:
                break
        answer += 1
    return answer

• 처음에는 리스트를 그대로 이용했다. enemy 배열을 순회했다.
• hap에 값을 더하고 빼면서 sum()을 이용하지 않으려고 했다.
• hap이 n을 넘고 사용할 수 있는 무적권이 남아있으면 
  • combs에서 가장 큰 값을 구해서 hap에서 빼고 k_cnt를 증가시킨 후에 combs 배열에서 가장 큰 값을 뺐다.
• 무적권이 남아있지 않으면 순회를 종료시켰다.
• 하지만 테스트 3, 6, 7, 9에서 시간초과가 발생했다.

 

정답 풀이

import heapq
def solution(n, k, enemy):
    answer = 0
    combs = []
    
    hap = 0
    k_cnt = 0
    for e in enemy:
        hap += e
        heapq.heappush(combs, (-e, e))
        
        if hap > n:
            if k_cnt < k:
                max_elem = heapq.heappop(combs)
                hap -= max_elem[1]
                k_cnt += 1
            else:
                break
        answer += 1
    return answer

로직에서 어차피 최대 값만 리스트에서 빼면 되는 것이라서, 최대 힙을 이용했다.

최대힙은 heapq.heappush(combs, (-e, e))로 -e를 기준으로 오름차순으로 힙이 정렬되는 것을 이용해야한다.

-e를 기준으로 오름차순으로 정렬되면 e를 기준으로는 내림차순이 되는 것이다.

자동으로 정렬이 되기때문에 max()를 사용할 일이 없어서 시간 복잡도가 줄어들게 된다.

 

리스트를 사용하면서 시간 초과가 발생하면 heapq를 쓰자!!