[프로그래머스] Swift - 이중 우선 순위 큐

 

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

프로그래머스

 

2023.07.03 기준 Level 3

 

 

알고리즘 초보가 공부용으로 풀고 기록하는 글입니다!
참고만 해주시고 더 좋은 풀이법이 있다면 알려주세요!

아이디어

우선 문제 이름부터 우선순위 큐를 사용하라고 힌트를 주고 있다..ㅎㅎ

문제 자체는 이해하기 쉽고 간단하다.

명령어가 I면 큐에 숫자 삽입

D 1 이면 큐에서 최댓값 삭제

D -1 이면 큐에서 최솟값 삭제이다.

 

결국 큐에서 최댓값, 최솟값을 삭제하는 로직이 필요한데 이러한 작업에 적합한 자료구조는 Heap이다.

우선순위 큐 또한 Heap을 이용해 구현하는 것이 일반적이다.

 

결국 이 문제는 Heap과 Priority Queue를 구현해서 최댓값과 최솟값을 제때 삭제하면 된다.

 

2023.07.03 기준으로 테스트 케이스가 매우 부족해보입니다.
Heap을 사용하지 않고 그냥 Array에 append한 뒤 필요할 때마다 매번 Sorting해도 테스트 케이스를 전부 통과하는 것 같다...
해당 방식대로라면 Level 1 수준이지만 그래도 연습 삼아 Heap을 사용해서 풀어보았고 역시 최선의 풀이가 아닙니다. 참고만 해주세요!

 

풀이

1. Heap 구현 (구글링 참고 및 수정해서 사용)

public struct Heap<T: Comparable> {
    private var elements: [T] = []
    var sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return self.elements.count <= 1
    }
    var peek: T? {
        if self.isEmpty { return nil }
        return self.elements.last!
    }
    var count: Int {
        return self.elements.count - 1
    }
    
    init(elements: [T] = [], sortFunction: @escaping (T, T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }
    
    private func leftChild(of index: Int) -> Int {
        return index * 2
    }
    
    private func rightChild(of index: Int) -> Int {
        return index * 2 + 1
    }
    
    private func parent(of index: Int) -> Int {
        return (index) / 2
    }
    
    private mutating func add(element: T) {
        self.elements.append(element)
    }
    
    mutating private func diveDown(from index: Int) {
        var higherPriority = index
        let leftIndex = self.leftChild(of: index)
        let rightIndex = self.rightChild(of: index)
        
        if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
            higherPriority = leftIndex
        }
        if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
            higherPriority = rightIndex
        }
        if higherPriority != index {
            self.elements.swapAt(higherPriority, index)
            self.diveDown(from: higherPriority)
        }
    }
    
    mutating private func swimUp(from index: Int) {
        var index = index
        while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
            self.elements.swapAt(index, self.parent(of: index))
            index = self.parent(of: index)
        }
    }
    
    mutating private func buildHeap() {
        for index in (1...(self.elements.count / 2)).reversed() {
            self.diveDown(from: index)
        }
    }
    
    mutating func reBuildHeap() {
        if elements.count > 1 {
            for index in (1...(self.elements.count / 2)).reversed() {
                self.diveDown(from: index)
            }
        }
    }
    
    mutating func insert(node: T) {
        if self.elements.isEmpty {
            self.elements.append(node)
        }
        self.elements.append(node)
        self.swimUp(from: self.elements.endIndex - 1)
    }
    
    @discardableResult
    mutating func remove() -> T? {
        if self.isEmpty { return nil }
        self.elements.swapAt(1, elements.endIndex - 1)
        let deleted = elements.removeLast()
        self.diveDown(from: 1)
        
        return deleted
    }
}

Heap은 워낙 자료가 다양하기 때문에 별다른 설명은 생략하겠습니다. 각자 익숙한 방식으로 구현해서 사용하면 됩니다!

다른 언어들에는 자체적으로 Heap을 제공하는 경우가 있는데 Swift에서는 거의 모든 자료구조를 직접 구현해야 합니다..ㅠ

그래서 이 문제의 난이도가 Level 3로 책정된 것 같습니다.

 

 

2. DoublePriorityQueue 구현

public struct DoublePriorityQueue<T: Comparable> {
    public enum HeapSortMode {
        case max
        case min
        
        var sortFunction: (T, T) -> Bool {
            switch self {
            case .max:
                return { $0 > $1 }
            case .min:
                return { $0 < $1 }
            }
        }
        
        mutating func toggle() {
            self = self == .max ? .min : .max
        }
    }
    
    var heap: Heap<T>
    
    var sortMode: HeapSortMode {
        didSet {
            self.heap.sortFunction = sortMode.sortFunction
            self.heap.reBuildHeap()
        }
    }
    
    var isEmpty: Bool {
        return self.heap.isEmpty
    }
    
    var count: Int {
        return self.heap.count
    }
    
    public init(elements: [T] = [], sortMode: HeapSortMode) {
        self.heap = Heap<T>(elements: elements, sortFunction: sortMode.sortFunction)
        self.sortMode = sortMode
    }
    
    mutating func enqueue(node: T) {
        self.heap.insert(node: node)
    }
    
    @discardableResult
    mutating func dequeueMax() -> T? {
        if self.heap.isEmpty { return nil }
        if self.sortMode == .max {
            return self.heap.remove()
        } else {
            self.sortMode.toggle()
            return self.heap.remove()
        }
    }
    
    @discardableResult
    mutating func dequeueMin() -> T? {
        if self.heap.isEmpty { return nil }
        if self.sortMode == .min {
            return self.heap.remove()
        } else {
            self.sortMode.toggle()
            return self.heap.remove()
        }
    }
    
    @discardableResult
    mutating func deqeue() -> T? {
        return self.heap.remove()
    }
}

이 문제의 핵심인 최댓값, 최솟값을 Dequeue하기 위해 Heap을 감싼 DoublePriorityQueue를 구현했습니다.

열거형으로 HeapSortMode 를 만들고 이 값을 .max, .min으로 토글하는 방식으로 Heap을 재정렬합니다.

max일 때는 최댓값 삭제, min일 때는 최솟값 삭제의 기능을 수행합니다.

 

3. Solution

fileprivate func solution(_ operations:[String]) -> [Int] {
    var result = [Int]()
    var queue = DoublePriorityQueue<Int>(sortMode: .max)
        
    for operation in operations {
        let operationArr = operation.split(separator: " ")
        let command = operationArr[0]
        let num = Int(operationArr[1])!
        
        switch (command, num) {
        case ("I", _):
            queue.enqueue(node: num)
        case ("D", 1):
            guard !queue.isEmpty else { break }
            queue.dequeueMax()
        case ("D", -1):
            guard !queue.isEmpty else { break }
            queue.dequeueMin()
        default:
            break
        }
    }
    
    if queue.isEmpty {
        return [0, 0]
    }
    
    let num = queue.deqeue()!
    result.append(num)
    
    if queue.isEmpty {
        return [num, num]
    }
        
    queue.sortMode.toggle()
    
    result.append(queue.deqeue()!)
    
    return result.sorted(by: >)
}

앞서 구현한 자료구조들을 바탕으로 solution 함수를 작성합니다.

충분히 추상화를 했기 때문에 solution 함수에서는 DoublePriorityQueue에 insert와 dequeue 작업만 수행하면됩니다.

 

함수 마지막 부분을 다시 보면

if queue.isEmpty {
        return [0, 0]
    }
    
let num = queue.deqeue()!
result.append(num)

if queue.isEmpty {
    return [num, num]
}

queue.sortMode.toggle()

result.append(queue.deqeue()!)

return result.sorted(by: >)

반복문을 마치고 qeueue가 비어있을 때 [0, 0]을 리턴합니다.

그리고 queue에 하나의 요소만 들어가 있을 경우를 구분하기 위해 한번 더 isEmpty를 검사하는 과정이 있는데 현재 프로그래머스 테스트케이스에서는 이 부분을 제거하고 곧바로 deqeue 2번으로 최댓값과 최솟값을 추출해도 통과한다고 합니다.

 

마무리

문제의 난이도 자체가 어렵거나 아이디어를 찾기 어려운 문제는 아니었습니다.

다만 Swift처럼 자료구조를 직접 구현해야 하는 언어에서는 꽤나 시간이 오래 걸릴 것 같습니다.

Heap과 PriorityQueue를 직접 구현해야 했기 때문에 자료구조 구현 연습에 도움이 되었던 그런 문제였습니다!!