99클럽 코테 스터디 12일차 TIL - 한계가 있는 풀이는 과감히 버리기

 

오늘 배운 점 

갑자기 떠오른 아이디어가 실마리가 될 수 있습니다.

빠를거라고 예상한 풀이가 예상외로 느릴 수 있습니다.

 

문제 - 디펜스 게임 문제링크

 

문제 설명

병사 n명으로 enemy([n1, n2, n3,...])을 상대해야 합니다. 1라운드에는 n1, 2라운드에는 n2... 이런식으로 각 라운드마다 상대해야 하는 적의 수가 enemy에 담겨 있습니다. 그리고 병사와 적은 1대1 교환비로 적을 처리할 수 있습니다.

k개의 무적권이 주어지는데, 무적권 1개를 사용하면 해당 라운드는 병사를 사용하지 않고 다음 라운드로 넘어갈 수 있습니다.

이때, 승리할 수 있는 최대 라운드를 구하면 됩니다.

제약 조건

n이 10억인 점은 문제가 되지 않지만 enemy가 100만이기 때문에 O(N)이나 O(NlogN) 과 같이 (N^2)을 넘지 않는 시간복잡도 풀이가 필요합니다.

 

풀이 과정

시행착오 1

이 문제를 처음 맞딱뜨렸을 때 DFS를 이용해서 풀어야겠다고 생각했습니다. enemy의 각 원소마다 무적권을 사용할지 말지 총 2가지 경우의 수가 있기 때문에 더이상 적을 처리할 수 없는 상태일 때를 종료조건으로 하는 DFS 재귀함수를 만들면 되겠다 라고 생각했습니다.

하지만 더 생각해보니 enemy의 크기가 100만이고, 각 원소마다 2가지 경우가 있기 때문에 총 경우의 수가 2^100만 이므로 절대 사용하면 안되는 시간복잡도였습니다.

 

시행착오 2

그래서, 한 번 순회하거나 이분탐색이 추가된 풀이를 생각해내려고 했습니다.

이후에 생각해낸 풀이는 이분탐색 기반의 파라메트릭 서치를 이용한 풀이였습니다.

우선, 서치 기준 값을 '이길 수 있는 라운드 값'으로 정한 후, start = 0, end = len(enemy)로 설정하였습니다.

이후, mid값을 구해 mid라운드까지 도달할 수 있는지를 체크한 후, 도달할 수 있으면 st를 증가시키고, 도달할 수 없으면 en를 감소하는 방향으로 구현하였습니다.

"""
    병사 n 10억
    무적권 k 50만
    적 enemy 100만

    파라매트릭 서치
    st, en : 클리어한 라운드 개수
    mid 구하고 0~mid 를 우큐에 넣고 제일 큰 K개 제거한 값이 n이하면 가능

"""
import heapq
def solution(n, k, enemy):
    if k >= len(enemy):
        return len(enemy)

    def solve(rnd):
        if k >= rnd:
            return True

        hq = []
        s = 0
        idx = 0
        while idx < rnd:
            s += enemy[idx]
            heapq.heappush(hq, -enemy[idx])
            idx += 1
        
        temp = k
        while temp > 0:
            s += heapq.heappop(hq)
            temp -= 1
        
        return s <= n
        
    st = 0
    en = len(enemy)
    ans = 0
    
    while st <= en:
        mid = (st + en) // 2
        res = solve(mid)

        if res:
            st = mid + 1 
            ans = max(ans, mid)
        else:
            en = mid - 1

    return ans

  

solve 함수에서 rnd 라운드에 도달할 수 있는지를 우선순위 큐를 이용해서 확인하는 것을 볼 수 있습니다.

rnd회 힙푸시를 하고, k만큼 힙팝을 하고 나서 적의 합과 병사를 비교합니다.

 

제출 결과

결과는 테스트 6에 대해 시간 초과하고 나머지는 전부 통과하였습니다.

이후 최대한 마이크로 최적화를 거쳤지만 시간 초과는 사라지지 않았습니다. 

 

시행착오 3

두번째 방법에서 제가 비효율적이라고 생각했던 부분은 solve 함수에서 rnd번 힙 푸시를 하고 다시 k 만큼 힙팝을 하는 과정이었습니다. 이분탐색을 거치며 이미 포함된 범위에 대해 계속 solve 가 진행될텐데 이 과정이 비효율적이라고 느껴졌습니다.

그래서 다른 방법을 생각했는데, 그건 기본적으로 enemy를 순회하며 더해주다가 n을 넘어설 경우 기존에 넘어갔던 라운드 중 가장 적이 많았던 라운드에 무적권을 사용하는 방법이었습니다.

가장 많았던 라운드를 알아내기 위해 매 순회마다 적의 수를 최대힙에 넣고, 필요시 최댓값을 pop 하는 방식으로 풀이했습니다.

import heapq
def solution(n, k, enemy):
    if k >= len(enemy):
        return len(enemy)
    
    hq = [] 
    _sum = 0
    cnts = 0
    for enm in enemy:
        heapq.heappush(hq, -enm)
        _sum += enm
        cnts += 1
        
        while k and _sum > n:
            k -= 1
            _sum += heapq.heappop(hq)
        
        if _sum > n:
            cnts -= 1
            break
    return cnts

 

매우 단순한 방법이고 코드도 훨씬 간단해진 모습을 확인할 수 있습니다.

위 풀이에서 주의해야 할 부분은, 더이상 추가할 수 없을 때의 cnts 처리입니다. for에 돌입하면 우선 cnts값이 +1 되면서 해당 라운드를 넘어갈 수 있다고 처리됩니다. 하지만 이후 무적권을 사용하고 나서도 sum > n 이라면 해당 라운드를 실패한 것이므로 다시 cnts -= 1 을 해주어야 합니다.

 

총평

아이디어가 떠오르면 생각외로 쉽게 풀리는 문제였습니다.

문제를 처음 봤을 때, 라운드가 진행되다가 n을 초과할 경우 지금까지 만난 적의 수 중 가장 큰 수를 무적권을 사용하고 싶다고 생각했는데, 이 생각이 현실화되게 하는 것이 우선순위 큐였다는 건 시행착오 3까지 오고나서 였습니다.