Swift) LeetCode(Easy) - Kth Largest Element In A Stream(Priority Queue)

leetcode.com/problems/kth-largest-element-in-a-stream/

Kth Largest Element in a Stream - LeetCode

안녕하세요 후르륵짭잡입니다!!

이번에는 우선순위 큐인 Heap 에 대해서 가져왔습니다.

사실 알고리즘을 잘 하지 못해서,,,,,, 구현 능력이 쫌 있을 뿐,,,, 생각 하는 능력이 떨어져서

이 문제는 여러번의 시간초과를 맞이 했습니다 ㅠㅠ.

 

일단 정답 코드를 보기 전에 Heap 코드를 보셔야 할 것 같습니다.

Heap은 우선 순위 큐로 Min / Max로 되어 있습니다.

기본적인 것은 다들 아실 거라 생각합니다 ㅎㅎㅎ

그런데 C++ 이나 JAVA , Python 같은 경우는 라이브러리가 있어서 

가져다가 쓰면 되지만 Swift는 그런거 없습니다.

그래서 복사 붙여넣기 해줘야합니다 ㅎㅎㅎ

 

** Heap 코드 ** 

struct Heap<Element>
{
    var elements : [Element]
    let priorityFunction : (Element, Element) -> Bool

    init(elements: [Element] = [], priorityFunction: @escaping (Element, Element) -> Bool) {
        self.elements = elements
        self.priorityFunction = priorityFunction
        buildHeap()
    }

    mutating func buildHeap() {
        for index in (0 ..< count / 2).reversed() {
            siftDown(elementAtIndex: index)
        }
    }

    var isEmpty : Bool { return elements.isEmpty }
    var count : Int { return elements.count }

    func peek() -> Element? {
        return elements.first
    }

    mutating func enqueue(_ element: Element) {
        elements.append(element)
        siftUp(elementAtIndex: count - 1)
    }

    mutating func siftUp(elementAtIndex index: Int) {
        let parent = parentIndex(of: index)
        guard !isRoot(index),
            isHigherPriority(at: index, than: parent)
            else { return }
        swapElement(at: index, with: parent)
        siftUp(elementAtIndex: parent)
    }

    mutating func dequeue() -> Element? {
        guard !isEmpty
            else { return nil }
        swapElement(at: 0, with: count - 1)
        let element = elements.removeLast()
        if !isEmpty {
            siftDown(elementAtIndex: 0)
        }
        return element
    }

    mutating func siftDown(elementAtIndex index: Int) {
        let childIndex = highestPriorityIndex(for: index)
        if index == childIndex {
            return
        }
        swapElement(at: index, with: childIndex)
        siftDown(elementAtIndex: childIndex)
    }

    // Helper functions

    func isRoot(_ index: Int) -> Bool {
        return index == 0
    }

    func leftChildIndex(of index: Int) -> Int {
        return (2 * index) + 1
    }

    func rightChildIndex(of index: Int) -> Int {
        return (2 * index) + 2
    }

    func parentIndex(of index: Int) -> Int {
        return (index - 1) / 2
    }

    func isHigherPriority(at firstIndex: Int, than secondIndex: Int) -> Bool {
        return priorityFunction(elements[firstIndex], elements[secondIndex])
    }

    func highestPriorityIndex(of parentIndex: Int, and childIndex: Int) -> Int {
        guard childIndex < count && isHigherPriority(at: childIndex, than: parentIndex)
            else { return parentIndex }
        return childIndex
    }

    func highestPriorityIndex(for parent: Int) -> Int {
        return highestPriorityIndex(of: highestPriorityIndex(of: parent, and: leftChildIndex(of: parent)), and: rightChildIndex(of: parent))
    }

    mutating func swapElement(at firstIndex: Int, with secondIndex: Int) {
        guard firstIndex != secondIndex
            else { return }
        elements.swapAt(firstIndex, secondIndex)
    }
}

이 코드를 보면 배열로 이루어져 있고,,, 먼가 복잡하구나를 느낄 겁니다 ㅎㅎㅎ

그런데, 사실 C++ 이나 어떤 언어든 Heap은 넣고 빼고만 알면 되져 ㅎㅎㅎㅎ

 let minHeap = Heap<Int>.init(elements: nums, priorityFunction: <)

만약에 minHeap을 구현하고 싶다면 Heap<제너릭 - 클래스 OR 프로토콜>.init(element : 배열 , priorityFuntion : < )을 해주면 됩니다.

maxHeap 같은 경우에는 minHeap 과 반대로 해주면 됩니다.

 maxHeap = Heap<Int>.init(elements: nums, priorityFunction: >)

만약에 배열이 없고 일반적인 prorityFuntion만 지정해준다고 한다면 아래 처럼 해주면 됩니다!

let minHeap = Heap<Int>(priorityFunction: >)

 

** Heap 삽입 과 삭제 그리고 TOP **

Heap 삽입 방법 :

minHeap.enqueue(val)

Heap 삭제 방법 :

let number = minHeap.dequeue()! // 옵셔널 타입 입니다

Heap ISEmpty:

minHeap.isEmpty //Bool형

Heap Top:

minHeap.peek()//옵셔널 타입

 

** 문제 해설 **

K 번째로 가장 큰 숫자를 리턴 해주는 겁니다.

 

** 많은 시도 **

정말 단순하게 접근 했습니다. 그래서 틀렸습니다 ㅎㅎㅎ

K 번 째로 가장 큰 숫자를 리턴하는 것이니 maxHeap으로 정렬해서 큰 숫자 부터 하나씩

빼버리면 된다고 생각 했는데,,, 결과는 시간초과,,,,

그래서 곰곰이 생각했습니다... (물론 다른 사람 코드도 참고 했습니다)

처음에는 Heap 코드에 문제가 있다 생각 했지만,,,, 문제가 있던건 제 머리 ㅎㅎㅎㅎ

 

- 배열로 풀었을 때 : 시간초과 -

class KthLargest {
    
    var minHeap : [Int] = []
    
    let largestK : Int
    
    init(_ k: Int, _ nums: [Int]) {
        
        minHeap = nums
        largestK = k

    }
    
    func add(_ val: Int) -> Int {
     
        minHeap.append(val)
        minHeap = minHeap.sorted(by: <)
        
        while minHeap.count > largestK {
            minHeap.remove(at: 0)
        }
        
        return minHeap[0]
    }
    
}

- MaxHeap으로 풀었을 때 : 시간초과 - 

class KthLargest {
    

    var numList : [Int] = []
    
    let largestK : Int
    
    init(_ k: Int, _ nums: [Int]) {
        
        numList = nums
        largestK = k

    }
    
    func add(_ val: Int) -> Int {
     
        numList.append(val)
        var maxHeap = Heap<Int>.init(elements: numList, priorityFunction: >)
        
        var temp = 1

        while temp < largestK {
            maxHeap.dequeue()!
            temp += 1
        }
        
        return maxHeap.dequeue()!
    }
    
}

 

이 두개를 봤을 때, 개인적으로 배열로 정렬하나 Heap으로 정렬하나 

시간 복잡도는 같을 줄 알았지만,,, 확실히 Heap이 빨랐습니다.

 

** 해결 방법 ** 

class KthLargest {

    var minHeap : Heap<Int>

    let largestK : Int

    init(_ k: Int, _ nums: [Int]) {

        minHeap = Heap<Int>.init(elements: nums, priorityFunction: <)
        largestK = k

    }

    func add(_ val: Int) -> Int {

        minHeap.enqueue(val)

        while minHeap.count > largestK {
            minHeap.dequeue()!
        }

        return minHeap.peek() ?? 0
    }

}

이 코드를 보면 무슨 느낌인지 알 겁니다.

만약 [4,5,8,2] 로 구성되어 있다면

k = 3 일 때, minHeap으로 정렬해서 [2,4,5,8]로 만듭니다.

그리고 작은 것을 Heap의 갯수가 K + 1 까지 삭제시키는 거죠.

이래도 되는 것이 어차피 max값은 살아 있고 맨 앞에 있는 것이 K번째 값일 겁니다.

 

이렇게 하니깐 시간초과가 발생하지 않고 문제를 해결 할 수 있었습니다.

 

Easy 였지만,,, Heap을 사용하지 않으면 풀수 없고 

또 너무 단순하게 얻어먹을려고 하면 안되는 문제였습니다... (반성 중)

모두 즐코코 하세욧!