[Swift]LeetCode373. 查找和最小的K对数字 | Find K Pairs with Smallest Sums

★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★
➤微信公众号：山青咏芝（shanqingyongzhi）
➤博客园地址：山青咏芝（https://www.cnblogs.com/strengthen/）
➤GitHub地址：https://github.com/strengthen/LeetCode
➤原文地址：https://www.cnblogs.com/strengthen/p/10278047.html 
➤如果链接不是山青咏芝的博客园地址，则可能是爬取作者的文章。
➤原文已修改更新！强烈建议点击原文地址阅读！支持作者！支持原创！
★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★★

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 
Explanation: The first 3 pairs are returned from the sequence: 
             [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: The first 2 pairs are returned from the sequence: 
             [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [1,3],[2,3]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]

---

给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

找到和最小的 k 对数字 (u1,v1), (u2,v2) ... (uk,vk)。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
     [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

---

52ms

class Solution {
    //Use priority queue
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        guard k > 0, nums1.count > 0, nums2.count > 0  else {
            return [[Int]]()
        }
        
        let pq = PriorityQueue<Pair>(priorityFunction:{(p1, p2) in
            return p1.sum < p2.sum
        })
        
        var result = [[Int]]()
        
        for j in 0..<nums2.count {
            let pair = Pair(i:0, j:j, sum:nums1[0] + nums2[j])
            pq.enqueue(element:pair)
        }
        
        
        
        
        for i in 0..<min(k, nums1.count * nums2.count) {
            if let cur = pq.dequeue() {
                result.append([nums1[cur.i], nums2[cur.j]])
                if cur.i == nums1.count - 1 {
                    continue
                }
                let pair = Pair(i:cur.i + 1, j:cur.j, sum:nums1[cur.i + 1] + nums2[cur.j])
                pq.enqueue(element:pair)
            }
        }
        
        return result
    }
}

class Pair {
    let i: Int
    let j: Int
    let sum: Int
    init(i:Int, j:Int, sum:Int) {
        self.i = i
        self.j = j
        self.sum = sum
    }
} 


public class PriorityQueue<Element> {
    var elements: [Element]
    var priorityFunction: (Element, Element) -> Bool
    var count: Int { return elements.count }
    
    init(priorityFunction:@escaping (Element, Element) -> Bool) {
        self.elements = [Element]()
        self.priorityFunction = priorityFunction
    }
    
    func isHigherPriority(at index:Int,than secondIndex:Int) -> Bool {
        return self.priorityFunction(elements[index], elements[secondIndex])
    }
    
    func enqueue(element:Element) {
        elements.append(element)
        siftUp(index: elements.count - 1)
    }
    
    func dequeue() -> Element? {
        if elements.count == 0 {
            return nil
        }
        
        elements.swapAt(0, elements.count - 1)
        let element = elements.removeLast()
        siftDown(index: 0)
        return element
    }
    
    func peek() -> Element? {
        return elements.last
    }
    
    func siftUp(index:Int) {
        if index == 0 {
            return
        }
        
        let parent = parentIndex(for: index)
        if isHigherPriority(at: index, than: parent) {
            elements.swapAt(index, parent)
            siftUp(index: parent)
        }
    }
    
    func siftDown(index:Int) {
        var highIndex = index
        let leftIndex = leftChildIndex(for: index)
        let rightIndex = rightChildIndex(for: index)
        if leftIndex < count && isHigherPriority(at: leftIndex, than: index) {
            highIndex = leftIndex
        }
        if rightIndex < count && isHigherPriority(at: rightIndex, than: highIndex) {
            highIndex = rightIndex
        }
        if highIndex == index {
            return
        } else {
            elements.swapAt(highIndex, index)
            siftDown(index: highIndex)
        }
    }
    
    func parentIndex(for index:Int) -> Int {
        return (index - 1)/2
    }
    
    func leftChildIndex(for index:Int) -> Int {
        return index * 2 + 1
    }
    
    func rightChildIndex(for index:Int) -> Int {
        return index * 2 + 2
    }
}

---

240ms

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
       
        var res = [[Int]]()
        let n1 = nums1.count 
        let n2 = nums2.count 
        guard n1 > 0 && n2 > 0 else {
            return res
        }
        var heap = Heap<Node>(type:. min)
        var visited = [[Bool]](repeating: [Bool](repeating: false, count: n2), count:n1)
        
        heap.add(Node(nums1[0] + nums2[0], [nums1[0], nums2[0]], 0, 0))
        visited[0][0] = true
        var count = k
        
        
        while count > 0 && heap.count > 0{
            var curr = heap.remove()!
            res.append(curr.arr)
            count -= 1
            if curr.x + 1 < n1 && !visited[curr.x + 1][curr.y]{
                let i = nums1[curr.x + 1]
                let j = nums2[curr.y]
                heap.add(Node(i + j, [i, j], curr.x + 1 , curr.y))
                visited[curr.x + 1][curr.y] = true
            }
            if curr.y + 1 < n2 && !visited[curr.x][curr.y + 1]{
                let i = nums1[curr.x]
                let j = nums2[curr.y + 1]
                heap.add(Node(i + j, [i, j], curr.x, curr.y + 1))
                visited[curr.x][curr.y + 1] = true
            }
        }
        return res
    }
    

struct Node: Comparable{
    var value : Int 
    var x:Int
    var y:Int
    var arr : [Int]
    init(_ n: Int, _ array:[Int],_ x:Int,_ y:Int){
        value = n
        arr = array
        self.x = x
        self.y = y
    }
    
    static func < (ls:Node ,rs: Node) -> Bool{
        return ls.value < rs.value
    }
    static func == (ls:Node ,rs: Node) -> Bool{
        return ls.value == rs.value
    }
}
    
    
struct Heap<E: Comparable> {
    enum HeapType { case max, min }

    // an array representing the binary tree
    var tree = [E]()

    // indicating if this heap is a max heap or min heap
    let type: HeapType

    var count: Int {
        return tree.count
    }

    init(type: HeapType) {
        self.type = type
    }

    mutating func add(_ element: E) {
        tree.append(element)
        var child = tree.count - 1
        var parent = (child - 1) / 2
        while child > 0 && !satify(tree[parent], tree[child]) {
            tree.swapAt(parent, child)
            child = parent
            parent = (child - 1) / 2
        }
        assert(isHeap(0), "borken heap: (tree)")
    }

    mutating func remove() -> E? {
        let rev = tree.first
        if tree.count > 0 {
            tree.swapAt(0, tree.count - 1)
            tree.removeLast()
            heapify(0)
        }
        assert(isHeap(0), "borken heap: (tree)")
        return rev
    }

    func peek() -> E? {
        return tree.first
    }

    mutating private func heapify(_ rootIndex: Int) {
        let count = tree.count
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2

        // no children
        if leftChild >= count { return }

        let targetChild = (rightChild >= count || satify(tree[leftChild], tree[rightChild])) ? leftChild : rightChild
        if (!satify(tree[rootIndex], tree[targetChild])) {
            tree.swapAt(rootIndex, targetChild)
            heapify(targetChild)
        }
    }

    private func satify(_ lhs: E, _ rhs: E) -> Bool {
        return (self.type == .min) ? (lhs <= rhs) : (lhs >= rhs)
    }

    // debugging helper methods
    private func isHeap(_ rootIndex: Int) -> Bool {
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2
        var valid = true
        if leftChild < tree.count {
            valid = satify(tree[rootIndex], tree[leftChild]) && isHeap(leftChild)
        }

        if !valid { return false }

        if rightChild < tree.count {
            valid = satify(tree[rootIndex], tree[rightChild]) && isHeap(rightChild)
        }
        return valid
    }
  }
}

---

844ms

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        
        guard nums1.count > 0 && nums2.count > 0 else {
            return []
        }
        
        var heap = PairHeap()
        var result = [[Int]]()
        var set = [(Int,Int)]()
        
        heap.addPair((0,0,nums1[0]+nums2[0]))
        for _ in 1...k {
            
            if heap.pairs.count > 0 {
                let cur = heap.poll()
                
                result.append([nums1[cur.0],nums2[cur.1]])
                
                var i = cur.0 + 1
                var j = cur.1
                
                
                if i < nums1.count && !(set.contains(where: { $0 == (i,j) })){ //if used hashmap, it will update eg i = 1 again n again
                    heap.addPair((i,j,nums1[i]+nums2[j]))
                    
                    set.append((i,j))
                }
                
                i = cur.0
                j = cur.1 + 1
                
                if j < nums2.count && !(set.contains(where: { $0 == (i,j) })){
                    heap.addPair((i,j,nums1[i]+nums2[j]))
                    set.append((i,j))
                    
                }
                //dump(heap.pairs)
            }
            //print(set)
        }
       
        return result
    }
    
    struct PairHeap {
    
    var pairs: [(Int,Int,Int)] = []
    
    func getLeftChildIndex(_ parentIndex: Int) -> Int{
        return 2*parentIndex + 1
    }
    
    func getRightChildIndex(_ parentIndex: Int) -> Int{
        return 2*parentIndex + 2
    }
    
    func getParentIndex(_ childIndex: Int) -> Int{
        return (childIndex-1)/2
    }
    
    // Return Item From Heap
    func getLeftChild(_ parenIndex: Int) -> (Int,Int,Int) {
        return pairs[getLeftChildIndex(parenIndex)]
    }
    
    func getRightChild(_ parenIndex: Int) -> (Int,Int,Int) {
        return pairs[getRightChildIndex(parenIndex)]
    }
    
    func getParent(_ childIndex: Int) -> (Int,Int,Int) {
         return pairs[getParentIndex(childIndex)]
    }
    
    // Boolean Check
    private func hasLeftChild(_ index: Int) -> Bool {
        return getLeftChildIndex(index) < pairs.count
    }
    private func hasRightChild(_ index: Int) -> Bool {
        return getRightChildIndex(index) < pairs.count
    }
    private func hasParent(_ index: Int) -> Bool {
        return getParentIndex(index) >= 0
    }
    
    mutating func addPair(_ item : (Int,Int,Int)) {
        pairs.append(item)
        heapifyUp()
    }
    
    mutating public func poll() -> (Int,Int,Int) {
        if pairs.count > 0 {
            let item = pairs[0]
            pairs[0] = pairs[pairs.count - 1]
            pairs.removeLast()
            heapifyDown()
            return item
            
        } else {
            fatalError()
        }
    }
    
    mutating func heapifyDown() {
        var index = 0
        
        while hasLeftChild(index) {
            var child = getLeftChildIndex(index)
            if hasRightChild(index) && getRightChild(index).2 < pairs[child].2 {
                child = getRightChildIndex(index)
            }
            
            if pairs[child].2 > pairs[index].2 {
                break
            } else {
                pairs.swapAt(index, child)
            }
            index = child
        }
        
    }
    
    mutating func heapifyUp() {
        var index = pairs.count - 1
        while hasParent(index) && getParent(index).2 > pairs[index].2 {
            pairs.swapAt(index, getParentIndex(index))
            index = getParentIndex(index)
        }   
    }    
  }
}

---

884ms

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        if nums1.count == 0 || nums2.count == 0 {
            return [[Int]]()
        }
        
        var rev = [[Int]]()
        var inHeap = Array(repeating: Array(repeating: false, count: nums2.count), count: nums1.count)
        var heap = Heap<Node>(type: .min)
        
        for i in 0 ..< nums2.count {
            let n = Node(pair: [nums1[0], nums2[i]], r: 0, c: i)
            heap.add(n)
        }
        
        while rev.count < k && heap.count > 0 {
            let n = heap.remove()!
            rev.append(n.pair)
            
            if n.r + 1 < nums1.count  {
                heap.add(Node(pair: [nums1[n.r + 1], nums2[n.c]], r: n.r + 1, c: n.c))
            }
        }
        
        return rev
    }
}

struct Node: Comparable {
    let pair: [Int]
    let r: Int
    let c: Int
    
    static func == (lhs: Node, rhs: Node) -> Bool {
        return (lhs.pair[0] + lhs.pair[1]) == (rhs.pair[0] + rhs.pair[1])
    }
    
    static func < (lhs: Node, rhs: Node) -> Bool {
        return (lhs.pair[0] + lhs.pair[1]) < (rhs.pair[0] + rhs.pair[1])
    }
}

struct Heap<E: Comparable> {
    enum HeapType { case max, min }

    // an array representing the binary tree
    var tree = [E]()

    // indicating if this heap is a max heap or min heap
    let type: HeapType

    var count: Int {
        return tree.count
    }

    init(type: HeapType) {
        self.type = type
    }

    mutating func add(_ element: E) {
        tree.append(element)
        var child = tree.count - 1
        var parent = (child - 1) / 2
        while child > 0 && !satify(tree[parent], tree[child]) {
            tree.swapAt(parent, child)
            child = parent
            parent = (child - 1) / 2
        }
        assert(isHeap(0), "borken heap: (tree)")
    }

    mutating func remove() -> E? {
        let rev = tree.first
        if tree.count > 0 {
            tree.swapAt(0, tree.count - 1)
            tree.removeLast()
            heapify(0)
        }
        assert(isHeap(0), "borken heap: (tree)")
        return rev
    }

    func peek() -> E? {
        return tree.first
    }

    mutating private func heapify(_ rootIndex: Int) {
        let count = tree.count
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2

        // no children
        if leftChild >= count { return }

        let targetChild = (rightChild >= count || satify(tree[leftChild], tree[rightChild])) ? leftChild : rightChild
        if (!satify(tree[rootIndex], tree[targetChild])) {
            tree.swapAt(rootIndex, targetChild)
            heapify(targetChild)
        }
    }

    private func satify(_ lhs: E, _ rhs: E) -> Bool {
        return (self.type == .min) ? (lhs <= rhs) : (lhs >= rhs)
    }

    // debugging helper methods
    private func isHeap(_ rootIndex: Int) -> Bool {
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2
        var valid = true
        if leftChild < tree.count {
            valid = satify(tree[rootIndex], tree[leftChild]) && isHeap(leftChild)
        }

        if !valid { return false }

        if rightChild < tree.count {
            valid = satify(tree[rootIndex], tree[rightChild]) && isHeap(rightChild)
        }
        return valid
    }
}

---

1512ms

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {

        var res = [[Int]]()
        for i in 0..<min(k, nums1.count) {
            for j in 0..<min(k, nums2.count) {
                res.append([nums1[i], nums2[j]])
            }
        }
        res.sort{$0[0]+$0[1] < $1[0]+$1[1]}
        return Array(res.prefix(k))
    }
}

---

2516ms

class Solution {
    func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
        var res:[[Int]] = [[Int]]()
        let number1:Int = min(nums1.count, k)
        let number2:Int = min(nums2.count, k)
        for i in 0..<number1
        {
            for j in 0..<number2
            {
                res.append([nums1[i], nums2[j]])
            }
        }        
        res.sort(by:{(a:[Int],b:[Int]) -> Bool in return a[0] + a[1] < b[0] + b[1] })        
        if res.count > k
        {
            var ans:[[Int]] = [[Int]]()
            for i in 0..<k
            {                
                ans.append(res[i])
            }
            return ans
        }
        else
        {
            return res
        }       
    }
}