[알고리즘] 프로그래머스 디펜스 게임

https://school.programmers.co.kr/learn/courses/30/lessons/142085

처음에 아무 생각 없이 enemy 리스트를 내림차순 정렬한 뒤 k개를 탐색한 대상들에서 무적권을 사용하면 될 것이라고 생각하였으나, 이 경우 무적권을 다 사용하지 못할 경우가 생기게 됨을 깨달았다.

일단 문제에서 주어진 enemy의 최대 범위가 enemy <= 1000000 이므로, O(n^2)으로는 풀 수 없고, O(nlogn)으로 풀어야함을 알 수 있다.

O(nlogn)이라고 하면 생각 나는 자료구조가 하나 있다. 바로 힙(heap) 이다. 힙을 이용해 푸는 아이디어는 다음과 같다.

1. enemy의 갯수보다 무적권의 개수가 더 많거나 같다면, 무조건 모든 라운드를 통과할 수 있다는 의미이기 때문에 그대로 enemy의 길이를 return 한다.
2. priority_queue의 크기에 k개의 enemy를 최소 힙으로 담아 둔다.

• 이는 곧 priority_queue에 무적권에 해당하는 enemy를 담아두고 꺼내지 않음으로써, n의 수를 차감해주지 않는 것을 의미한다.

3. priority_queue의 갯수가 k개를 넘어선다면, 하나씩 pop해서 n에서 차감해준다. 이때 pop한 값이 n보다 크다면, 더 이상 라운드를 진행할 수 없음을 의미하므로 라운드 수를 return 해준다.

정답 코드

import heapq
def solution(n, k, enemy):
    if k >= len(enemy):
        return len(enemy)
    priority_queue = []
    
    for i in range(len(enemy)):
        heapq.heappush(priority_queue, enemy[i])
        if len(priority_queue) > k:
            last = heapq.heappop(priority_queue)
            if last > n:
                return i
            n -= last
    return len(enemy)

어떻게 이런 코드가 논리적으로 가능할까?

답은 무적권을 고르는 기준에 있다. 나는 이 기준을 떠올리지 못해서 문제 풀이 방법을 생각해내지 못했다.
무적권은 현재 지나온 라운드 중에서 가장 큰 k개의 값들이 무적권의 혜택을 받을 수 있는 값들이 된다. 따라서, 이 값들을 힙에 최소힙 기준으로 넣어주고 k개 초과가 되었을때 pop해주게 되면, 자연스럽게 현재 지나온 값들 중 큰 k개의 값들은 그대로 힙에 남아있게 되고 나머지 값들만 pop되어서 n을 차감하는 연산에 사용될 수 있게 된다.

아직 힙이라는 자료 구조를 조금 더 유연하게 생각하는 사고력이 모자르다는것을 느낄 수 있는 문제였다.