코딩테스트 연습 - 이중우선순위큐

문제 풀이 아이디어

중간에 입력이 주어질 때마다 정렬을 해야 하는 문제이므로 힙으로 구현해야 하는 문제입니다. 다만 문제점이 힙은 최소힙 아니면 최대힙으로 정렬을 내림차순 혹은 오름차순 중에 하나로만 구현할 수 있는데 이 문제의 경우 최소, 최대 값을 오가면서 구해야 하는 문제입니다.

힙을 하나 사용하는 경우

일단 최소힙만 사용하면서 수들을 push 합니다. 그리고 나서 최솟값을 pop해야 할 때는 그냥 pop하고 최댓값을 pop해야 할 때는 최소힙에서 모든 수를 pop하면서 마지막에 pop된 수 (= 최댓값)을 남기고 나머지는 도로 최소힙에 넣습니다.

시간 복잡도를 계산해보면 최솟값을 구할 때는 O(1)이고 최댓값을 구할 때는 모든 수를 pop할 때 O(n)이고 최댓값을 제외한 모든 수를 다시 push할 때 O(nlogn)입니다. 따라서 O(n) + O(nlogn) = O(nlogn)이 되겠네요.

힙을 두 개 사용하는 경우

이번에는 최소힙과 최대힙을 동시에 사용해봅시다. 최솟값을 구해야 할 때는 최소힙에서 최댓값을 구할 때는 최대힙에서 pop하면 됩니다. 하지만 하나의 힙에서만 pop하는 것이 아니라 반대편 heap에서도 pop해야 합니다. 힙을 하나 사용하는 경우와 마찬가지로 모든 수를 pop하고 나서 마지막 수를 제외한 수들은 다시 push 해주어야 합니다.

시간복잡도는 하나를 사용하는 경우와 마찬가지입니다. 하지만 이 경우에는 최솟값, 최댓값을 사용하는 경우 모두 O(nlogn)이 되겠네요.

정리

뭘 사용해도 최악의 케이스를 가정하는 Big-O 시간복잡도 기준으로는 O(nlogn)으로 동일합니다. 저는 최소힙 하나만 사용해서 구현해보겠습니다.

코드

// 최소힙 구현
import Foundation

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

    var isEmpty: Bool {
        return heap.count <= 1 ? true : false
    }
    
    
    init() {}

    init(_ element: T) {
        heap.append(element)
        heap.append(element)
    }
    
    mutating func insert(_ element: T) {
        
        if heap.isEmpty {
            heap.append(element)
            heap.append(element)
            return
        }
        
        heap.append(element)
        
        func isMoveUp(_ insertIndex: Int) -> Bool {
            if insertIndex <= 1 {
                return false
            }
            
            let parentIndex = insertIndex / 2
            return heap[insertIndex] < heap[parentIndex] ? true : false
        }
        
        var insertIndex = heap.count - 1
        
        while isMoveUp(insertIndex) {
            let parentIndex = insertIndex / 2
            heap.swapAt(insertIndex, parentIndex)
            insertIndex = parentIndex
        }
    }
    

    enum moveDownDirection { case left, right, none }
    

    mutating func pop() -> T? {
        
        if heap.count <= 1 {
            return nil
        }
        
        let returnData = heap[1]
        heap.swapAt(1, heap.count - 1)
        heap.removeLast()
        
        func moveDown(_ poppedIndex: Int) -> moveDownDirection {
            
            let leftChildIndex = poppedIndex * 2
            let rightChildIndex = leftChildIndex + 1
            
            if leftChildIndex >= heap.count {
                return .none
            }
            
            if rightChildIndex >= heap.count {
                return heap[leftChildIndex] < heap[poppedIndex] ? .left : .none
            }
            
            if (heap[leftChildIndex] > heap[poppedIndex]) && (heap[rightChildIndex] > heap[poppedIndex]) {
                return .none
            }
            
            if (heap[leftChildIndex] < heap[poppedIndex]) && (heap[rightChildIndex] < heap[poppedIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            if (heap[leftChildIndex] < heap[poppedIndex]) || (heap[rightChildIndex] < heap[poppedIndex]) {
                return heap[leftChildIndex] < heap[rightChildIndex] ? .left : .right
            }
            
            return .none
        }
        
        var poppedIndex = 1
        while true {
            switch moveDown(poppedIndex) {
            case .none:
                return returnData
            case .left:
                let leftChildIndex = poppedIndex * 2
                heap.swapAt(poppedIndex, leftChildIndex)
                poppedIndex = leftChildIndex
            case .right:
                let rightChildIndex = (poppedIndex * 2) + 1
                heap.swapAt(poppedIndex, rightChildIndex)
                poppedIndex = rightChildIndex
            }
        }
    }
}

// 오퍼레이션 타입 정의
enum Operation {
    case insert(value: Int)
    case minPop
    case maxPop
}

// 문자열을 Operation 인스턴스로 바꾸어주는 함수
func parseOperation(_ string: String) -> Operation {
    let array = string.split(separator: " ")
    switch array[0] {
    case "I": return .insert(value: Int(array[1])!)
    case "D":
        if array[1] == "1" { return .maxPop }
        else { return .minPop }
    default: fatalError()
    }
}

func solution(_ operations:[String]) -> [Int] {
    // 최댓값을 pop하는 방법
    func popMax() -> Int? {
        // 힙이 비었다면 nil을 리턴
        guard !heap.isEmpty else { return nil }
        
        // 전부 빼서 마지막만 return 하기
        var temp = [Int]()
        while !heap.isEmpty {
            temp.append(heap.pop()!)
        }
        let max = temp.popLast()!
        while !temp.isEmpty {
            heap.insert(temp.popLast()!)
        }
        return max
    }
    
    // 문자열로 들어온 명령어 파싱
    let parsed = operations.map { parseOperation($0) }
    
    // 힙 선언
    var heap = MinHeap<Int>()
    
    // 명령어 실행
    for oper in parsed {
        switch oper {
        case .insert(value: let value):
            heap.insert(value)
        case .minPop:
            _ = heap.pop()
        case .maxPop:
            _ = popMax()
        }
    }
    
    // 실행 마친 후 최소값과 최대값 구하기
    let min = heap.pop() ?? 0
    let max = popMax() ?? 0
    
    return [max, min]
}