[Heap] 우선순위 큐 - LeetCode Swift

문제 : https://leetcode.com/problems/find-k-pairs-with-smallest-sums/

Find K Pairs with Smallest Sums - LeetCode

 

우선 이 문제는 우선순위큐를 먼저 이해하고 직접 코드로 구현할 수 있는가? 묻는 문제이다.

우선순위큐란 큐에 들어온 순서에 상관없이 우선순위가 높은 데이터가 먼저 나오는 것이다.

 

코드

//
//  main.swift
//  FindKPairsWithSmallestSums
//
//  Created by 오킹 on 2021/08/11.
//

import Foundation

class Solution {
    struct Pair {
        var left: Int
        var right: Int
        var index: Int
        var sum: Int
        
        init(_ left: Int, _ right: Int, _ index: Int) {
            self.left = left
            self.right = right
            self.sum = left + right
            self.index = index
        }
    }
    
    
    class PriorityQueue {
        private var arr = [Pair]()
        
        var arrLog: [[Int]] {
            return arr.map({[$0.left, $0.right, $0.index]})
        }
        
        var pull: Pair? {
            return arr.first
        }
        
        var isEmpty: Bool {
            return arr.isEmpty
        }
        
        func pop() -> Pair {
            let temp = arr.first!
            arr[0] = arr.last!
            arr.removeLast()
            heapifyTowDown(0)
            return temp
        }
        
        func insert(_ pair: Pair) {
            arr.append(pair)
            heapifyBottomUp(arr.count - 1)
        }
        
        private func heapifyBottomUp(_ index: Int) {
            if index == 0 {
                return
            }
            
            if arr[parentIndex(index)].sum > arr[index].sum {
                swap(parentIndex(index), index)
                heapifyBottomUp(parentIndex(index))
            }
        }
        
        private func heapifyTowDown(_ index: Int) {
            if index >= arr.count {
                return
            }
            
            let leftChildIndex = leftChild(index)
            let rightChildIndex = rightChild(index)
            
            if rightChildIndex < arr.count {
                if arr[leftChildIndex].sum <= arr[rightChildIndex].sum && arr[leftChildIndex].sum < arr[index].sum {
                    swap(leftChildIndex, index)
                    heapifyTowDown(leftChildIndex)
                } else if arr[rightChildIndex].sum < arr[index].sum {
                    swap(rightChildIndex, index)
                    heapifyTowDown(rightChildIndex)
                }
            } else if leftChildIndex < arr.count && arr[leftChildIndex].sum < arr[index].sum {
                swap(leftChildIndex, index)
                heapifyTowDown(leftChildIndex)
            }
        }
        
        private func swap(_ indexOne: Int, _ indexTwo: Int) {
            let temp = arr[indexOne]
            arr[indexOne] = arr[indexTwo]
            arr[indexTwo] = temp
        }
        
        private func parentIndex(_ index: Int) -> Int {
            return (index - 1)/2
        }
        
        private func leftChild(_ index: Int) -> Int {
            return 2 * index + 1
        }
        
        private func rightChild(_ index: Int) -> Int {
            return 2 * index + 2
        }
    }
    
    
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        var res = [[Int]]()
        
        if nums1.isEmpty || nums2.isEmpty {
            return []
        }
        
        var queue = PriorityQueue()
        for i in 0..<nums1.count {
            var pair = Pair(nums1[i], nums2[0], 0)
            queue.insert(pair)
        }
        
        var k = k
        while k > 0 && !queue.isEmpty {
            k -= 1
            let current = queue.pop()
            res.append([current.left, current.right])
            if current.index == nums2.count - 1 {continue}
            queue.insert(Pair(current.left, nums2[current.index + 1], current.index + 1))
        }
        
        return res
    }
}

 

아직 완벽하게 이해하진 못했지만 다른 문제를 만난다면 보면서라도 이 방법을 떠올려서 풀어봐야겠다.

코드를 이해하기 위한 지식 : https://blog.weirdx.io/post/3122

출처 : https://leetcode.com/problems/find-k-pairs-with-smallest-sums/discuss/434185/Swift-Priority-queue100-faster.-Custom-priority-queue-implementation.