문제

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

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

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

문제 풀러 가기


풀이


우선순위 큐 사용

  1. 문제 접근
    • 무적권을 가장 적합한 시기에 사용하기 위해 등장한 적의 수 중 가장 많은 라운드에 무적권을 우선적으로 사용합니다.
  2. 풀이 절차
    • 입력 값 확인 및 초기 설정
      • 주어진 라운드 수가 무적권 사용 횟수 k보다 작거나 같다면 모든 라운드를 무적권으로 막을 수 있으므로 즉시 종료합니다. 무적권을 사용할 라운드 관리를 위해 최소 힙 (Min-Heap) 을 생성합니다.
  • 라운드 순회 및 처리
    • 각 라운드의 적 수를 힙에 추가합니다.
    • 힙의 크기가 k를 초과하면, 병사 소모를 줄이기 위해 가장 적은 적 수를 가진 라운드부터 병사로 막습니다.
    • 병사가 부족하면 게임을 종료하고, 현재 라운드 번호를 반환합니다.
  • 최대 라운드 계산
    • 라운드 순회를 모두 완료하면, 모든 라운드를 막을 수 있으므로 총 라운드 수를 반환합니다.



코드

import heapq

def solution(n, k, enemy):
    # 전체 라운드 수
    stage = len(enemy)
    
    # 무적권으로 모든 라운드 막을 수 있는 경우
    if k >= stage:
        return stage
    
    # 최소 힙 초기화
    q = []
    
    for i in range(stage):
        # 현재 라운드의 적 수를 힙에 추가
        heapq.heappush(q, enemy[i])
        
        # 힙 크기가 k를 초과하면, 가장 작은 적 수만큼 병사로 막기
        if len(q) > k:
            last = heapq.heappop(q)
            if last > n:  # 병사가 부족한 경우
                return i
            n -= last  # 병사 소모
    
    # 모든 라운드를 막은 경우
    return stage