[Algorithm] 프로그래머스 - 디펜스 게임 (Swift) (Heap, Parametric Search)

문제 설명

프로그래머스

문제가 아주 심플한데, n명의 병사, k번의 무적권(?)을 가지고 enemy 배열이 주어졌을 때 최대 몇 라운드까지 갈 수 있는지를 구하는 문제였다.

 

문제 풀이 과정

문제 접근은 아래와 같다. 우선 무적권를 초반에 몰아서 쓴다. 그리고 그것을 stack에 넣어서 stack의 요소 중에 가장 적은 요소보다 큰 enemy가 나왔을 때 무적권을 취소하고, 그 enemy에 무적권을 쓰고, 그 제일 작은 요소에 병사를 사용하는 것으로 문제를 해결했다. 일종의 그리디 알고리즘 느낌이랄까...? 그때그때 최적을 선택하는 방식으로 진행했다.

import Foundation

func solution(_ n:Int, _ k:Int, _ enemy:[Int]) -> Int {
    var n = n    
    var k = k
    var stack = enemy[0..<k]
    stack.sort(by: >)
    
    if k >= enemy.count { return enemy.count }
    
    while k < enemy.count {
        let poped = stack.removeLast()
       if enemy[k] > poped {
           n -= poped           
           stack.append(enemy[k])
           stack.sort(by: >)
       } else {
           n -= enemy[k]
           stack.append(poped)
       }
        
        if n <= 0 { break }
        k += 1
    }
    
    return k
}

 

결과는 시간 초과 오류가 나왔다. 시간복잡도가 O(NlogN)인 sort를 사용했기 때문에 (N의 최대가 10억) 시간 초과가 나는 것은 당연했다.

 

Heap 사용

자료를 효과적으로 넣고, 효과적으로 추출하기 위해서는 O(logN)의 시간복잡도를 가지고 있는 Heap이 필요하다는 것을 느꼈다. 힙구조를 직접 구현하여 stack을 heap으로만 바꿔주었다.

import Foundation

func solution(_ n:Int, _ k:Int, _ enemy:[Int]) -> Int {
    var n = n
    var k = k
    var heap = MinHeap<Int>()
    
    if k >= enemy.count { return enemy.count }
    
    for i in 0..<k {
        heap.push(enemy[i])
    }
    
    while k < enemy.count {
        let poped = heap.pop()!
        if enemy[k] > poped {
            n -= poped
            heap.push(enemy[k])
        } else {
            n -= enemy[k]
            heap.push(poped)
        }

        if n < 0 { break }
        k += 1
    }
    
    return k
}

struct MinHeap <T: Comparable> {
    var heap: [T] = []

    func isEmpty() -> Bool {
        heap.count <= 1
    }

    mutating func push(_ data: T) {
        if isEmpty() { heap.append(data) }
        var index = heap.count
        heap.append(data)

        while index > 1 {
            let parent = heap[index/2]
            if parent < data { break }
            heap[index] = parent
            index /= 2
        }
        heap[index] = data
    }

    mutating func pop() -> T? {
        if isEmpty() { return nil }
        let item = heap[1]
        let data = heap.removeLast()
        let size = heap.count - 1
        
        if isEmpty() { return item }
        heap[1] = data
        
        var (parent, child) = (1, 2)
        while child <= size {
            if child < size && heap[child] > heap[child+1] {
                child += 1
            }
            if heap[parent] < heap[child] { break }
            heap.swapAt(parent, child)
            parent = child
            child = parent * 2
        }
        
        return item
    }
}

 

결과는 훨씬 개선된 결과가 나왔다.

 

Parametric Search

문제를 풀고나서 파라메트릭 서치를 사용하여 문제를 풀 수 있다는 것을 알게 되어서 이것도 시도해 보았다. 코드는 좀 더 짧아졌으나... 역시 sorted가 들어갔기 때문에 시간초과가 나게 되었지만 파라메트릭 서치를 사용할 수 있는 기회가 있어서 좋았다.

import Foundation

func solution(_ n:Int, _ k:Int, _ enemy:[Int]) -> Int {
    if k >= enemy.count { return enemy.count }
    var (left, right) = (0, enemy.count)
    
    func defence(_ round: Int) -> Bool {
        var (n, k) = (n, k)
        if k >= round { return true }
        let stages = enemy[0..<round].sorted(by: >)
                
        for i in k..<round {
            n -= stages[i]
            if n < 0 { return false }
        }
        
        return n >= 0 ? true : false
    }
    
    while left <= right {
        let mid = (left + right)/2
        if defence(mid) {
            left = mid + 1
        } else {
            right = mid - 1
        }
    }
    
    return left-1
}