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

문제 요약

병사 n명으로 순서대로 등장하는 적을 막는다.

• 한 라운드를 일반적으로 막으면 적 수만큼 병사가 줄어든다.
• 무적권은 최대 k번 사용할 수 있고, 사용한 라운드에서는 병사가 줄지 않는다.

최대로 막을 수 있는 라운드 수를 구한다.

핵심 아이디어

어떤 시점까지 막은 라운드들 중에서 무적권은 적 수가 가장 많은 k개 라운드에 쓰는 것이 항상 최적이다.

적 수가 작은 라운드에 무적권을 쓰고 큰 라운드에 병사를 쓰는 경우가 있다면, 두 라운드의 무적권 사용 여부를 바꾸면 병사 소모는 줄어들거나 같아진다.

따라서 각 라운드에서 다음 상태를 유지한다.

• 최소 힙: 지금까지 등장한 공격 중 무적권을 사용할 k개의 큰 공격
• 병사: 무적권을 쓰지 않는 나머지 공격을 막는 데 사용

힙 크기가 k를 넘으면, 힙에서 가장 작은 공격을 꺼내 병사로 막는다.

풀이 과정

1. 적 수를 최소 힙에 넣는다.
2. 힙의 크기가 k보다 크면, 가장 작은 적 수를 꺼내 병사에서 뺀다.
3. 병사가 음수가 되면 현재 라운드를 막을 수 없으므로, 현재 인덱스를 반환한다.
4. 모든 라운드를 처리하면 전체 라운드 수를 반환한다.

Python 코드

import heapq


def solution(n, k, enemy):
    invincible_rounds = []

    for round_index, enemy_count in enumerate(enemy):
        # 현재까지는 이 라운드에도 무적권을 쓴다고 가정한다.
        heapq.heappush(invincible_rounds, enemy_count)

        # 무적권 수를 넘으면 가장 작은 공격은 병사로 막는다.
        if len(invincible_rounds) > k:
            n -= heapq.heappop(invincible_rounds)

        if n < 0:
            return round_index

    return len(enemy)

예시

n = 7, k = 3, enemy = [4, 2, 4, 5, 3, 3, 1]인 경우를 보자.

5라운드까지 처리한 뒤 무적권을 사용할 공격은 가장 큰 5, 4, 4다.

무적권 사용: 5, 4, 4
병사로 처리: 2 + 3 = 5
남은 병사: 7 - 5 = 2

6라운드의 적 수 3까지 포함하면, 무적권을 제외하고 병사로 막아야 하는 적 수의 합이 10이 된다. 병사가 부족하므로 5라운드까지만 막을 수 있다.

시간 복잡도

N을 전체 라운드 수라고 하자.

힙에는 최대 k개의 원소만 유지한다.

• 힙 삽입과 삭제: O(log k)
• 전체 라운드 처리: O(N log k)
• 공간 복잡도: O(k)

정리

매 순간 무적권을 가장 큰 공격들에 배정한다고 생각하면 된다. 최소 힙에서 가장 작은 값을 병사로 처리하면서, 지금까지의 최적 무적권 배정을 계속 유지할 수 있다.