[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

  1 class Solution {
  2     //Use priority queue
  3     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
  4         guard k > 0, nums1.count > 0, nums2.count > 0  else {
  5             return [[Int]]()
  6         }
  7         
  8         let pq = PriorityQueue(priorityFunction:{(p1, p2) in
  9             return p1.sum < p2.sum
 10         })
 11         
 12         var result = [[Int]]()
 13         
 14         for j in 0..<nums2.count {
 15             let pair = Pair(i:0, j:j, sum:nums1[0] + nums2[j])
 16             pq.enqueue(element:pair)
 17         }
 18         
 19         
 20         
 21         
 22         for i in 0.. nums2.count) {
 23             if let cur = pq.dequeue() {
 24                 result.append([nums1[cur.i], nums2[cur.j]])
 25                 if cur.i == nums1.count - 1 {
 26                     continue
 27                 }
 28                 let pair = Pair(i:cur.i + 1, j:cur.j, sum:nums1[cur.i + 1] + nums2[cur.j])
 29                 pq.enqueue(element:pair)
 30             }
 31         }
 32         
 33         return result
 34     }
 35 }
 36 
 37 class Pair {
 38     let i: Int
 39     let j: Int
 40     let sum: Int
 41     init(i:Int, j:Int, sum:Int) {
 42         self.i = i
 43         self.j = j
 44         self.sum = sum
 45     }
 46 } 
 47 
 48 
 49 public class PriorityQueue {
 50     var elements: [Element]
 51     var priorityFunction: (Element, Element) -> Bool
 52     var count: Int { return elements.count }
 53     
 54     init(priorityFunction:@escaping (Element, Element) -> Bool) {
 55         self.elements = [Element]()
 56         self.priorityFunction = priorityFunction
 57     }
 58     
 59     func isHigherPriority(at index:Int,than secondIndex:Int) -> Bool {
 60         return self.priorityFunction(elements[index], elements[secondIndex])
 61     }
 62     
 63     func enqueue(element:Element) {
 64         elements.append(element)
 65         siftUp(index: elements.count - 1)
 66     }
 67     
 68     func dequeue() -> Element? {
 69         if elements.count == 0 {
 70             return nil
 71         }
 72         
 73         elements.swapAt(0, elements.count - 1)
 74         let element = elements.removeLast()
 75         siftDown(index: 0)
 76         return element
 77     }
 78     
 79     func peek() -> Element? {
 80         return elements.last
 81     }
 82     
 83     func siftUp(index:Int) {
 84         if index == 0 {
 85             return
 86         }
 87         
 88         let parent = parentIndex(for: index)
 89         if isHigherPriority(at: index, than: parent) {
 90             elements.swapAt(index, parent)
 91             siftUp(index: parent)
 92         }
 93     }
 94     
 95     func siftDown(index:Int) {
 96         var highIndex = index
 97         let leftIndex = leftChildIndex(for: index)
 98         let rightIndex = rightChildIndex(for: index)
 99         if leftIndex < count && isHigherPriority(at: leftIndex, than: index) {
100             highIndex = leftIndex
101         }
102         if rightIndex < count && isHigherPriority(at: rightIndex, than: highIndex) {
103             highIndex = rightIndex
104         }
105         if highIndex == index {
106             return
107         } else {
108             elements.swapAt(highIndex, index)
109             siftDown(index: highIndex)
110         }
111     }
112     
113     func parentIndex(for index:Int) -> Int {
114         return (index - 1)/2
115     }
116     
117     func leftChildIndex(for index:Int) -> Int {
118         return index * 2 + 1
119     }
120     
121     func rightChildIndex(for index:Int) -> Int {
122         return index * 2 + 2
123     }
124 }

---

240ms

  1 class Solution {
  2     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
  3        
  4         var res = [[Int]]()
  5         let n1 = nums1.count 
  6         let n2 = nums2.count 
  7         guard n1 > 0 && n2 > 0 else {
  8             return res
  9         }
 10         var heap = Heap(type:. min)
 11         var visited = [[Bool]](repeating: [Bool](repeating: false, count: n2), count:n1)
 12         
 13         heap.add(Node(nums1[0] + nums2[0], [nums1[0], nums2[0]], 0, 0))
 14         visited[0][0] = true
 15         var count = k
 16         
 17         
 18         while count > 0 && heap.count > 0{
 19             var curr = heap.remove()!
 20             res.append(curr.arr)
 21             count -= 1
 22             if curr.x + 1 < n1 && !visited[curr.x + 1][curr.y]{
 23                 let i = nums1[curr.x + 1]
 24                 let j = nums2[curr.y]
 25                 heap.add(Node(i + j, [i, j], curr.x + 1 , curr.y))
 26                 visited[curr.x + 1][curr.y] = true
 27             }
 28             if curr.y + 1 < n2 && !visited[curr.x][curr.y + 1]{
 29                 let i = nums1[curr.x]
 30                 let j = nums2[curr.y + 1]
 31                 heap.add(Node(i + j, [i, j], curr.x, curr.y + 1))
 32                 visited[curr.x][curr.y + 1] = true
 33             }
 34         }
 35         return res
 36     }
 37     
 38 
 39 struct Node: Comparable{
 40     var value : Int 
 41     var x:Int
 42     var y:Int
 43     var arr : [Int]
 44     init(_ n: Int, _ array:[Int],_ x:Int,_ y:Int){
 45         value = n
 46         arr = array
 47         self.x = x
 48         self.y = y
 49     }
 50     
 51     static func < (ls:Node ,rs: Node) -> Bool{
 52         return ls.value < rs.value
 53     }
 54     static func == (ls:Node ,rs: Node) -> Bool{
 55         return ls.value == rs.value
 56     }
 57 }
 58     
 59     
 60 struct Heap {
 61     enum HeapType { case max, min }
 62 
 63     // an array representing the binary tree
 64     var tree = [E]()
 65 
 66     // indicating if this heap is a max heap or min heap
 67     let type: HeapType
 68 
 69     var count: Int {
 70         return tree.count
 71     }
 72 
 73     init(type: HeapType) {
 74         self.type = type
 75     }
 76 
 77     mutating func add(_ element: E) {
 78         tree.append(element)
 79         var child = tree.count - 1
 80         var parent = (child - 1) / 2
 81         while child > 0 && !satify(tree[parent], tree[child]) {
 82             tree.swapAt(parent, child)
 83             child = parent
 84             parent = (child - 1) / 2
 85         }
 86         assert(isHeap(0), "borken heap: \(tree)")
 87     }
 88 
 89     mutating func remove() -> E? {
 90         let rev = tree.first
 91         if tree.count > 0 {
 92             tree.swapAt(0, tree.count - 1)
 93             tree.removeLast()
 94             heapify(0)
 95         }
 96         assert(isHeap(0), "borken heap: \(tree)")
 97         return rev
 98     }
 99 
100     func peek() -> E? {
101         return tree.first
102     }
103 
104     mutating private func heapify(_ rootIndex: Int) {
105         let count = tree.count
106         let leftChild = 2 * rootIndex + 1
107         let rightChild = 2 * rootIndex + 2
108 
109         // no children
110         if leftChild >= count { return }
111 
112         let targetChild = (rightChild >= count || satify(tree[leftChild], tree[rightChild])) ? leftChild : rightChild
113         if (!satify(tree[rootIndex], tree[targetChild])) {
114             tree.swapAt(rootIndex, targetChild)
115             heapify(targetChild)
116         }
117     }
118 
119     private func satify(_ lhs: E, _ rhs: E) -> Bool {
120         return (self.type == .min) ? (lhs <= rhs) : (lhs >= rhs)
121     }
122 
123     // debugging helper methods
124     private func isHeap(_ rootIndex: Int) -> Bool {
125         let leftChild = 2 * rootIndex + 1
126         let rightChild = 2 * rootIndex + 2
127         var valid = true
128         if leftChild < tree.count {
129             valid = satify(tree[rootIndex], tree[leftChild]) && isHeap(leftChild)
130         }
131 
132         if !valid { return false }
133 
134         if rightChild < tree.count {
135             valid = satify(tree[rootIndex], tree[rightChild]) && isHeap(rightChild)
136         }
137         return valid
138     }
139   }
140 }

---

844ms

  1 class Solution {
  2     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
  3         
  4         guard nums1.count > 0 && nums2.count > 0 else {
  5             return []
  6         }
  7         
  8         var heap = PairHeap()
  9         var result = [[Int]]()
 10         var set = [(Int,Int)]()
 11         
 12         heap.addPair((0,0,nums1[0]+nums2[0]))
 13         for _ in 1...k {
 14             
 15             if heap.pairs.count > 0 {
 16                 let cur = heap.poll()
 17                 
 18                 result.append([nums1[cur.0],nums2[cur.1]])
 19                 
 20                 var i = cur.0 + 1
 21                 var j = cur.1
 22                 
 23                 
 24                 if i < nums1.count && !(set.contains(where: { $0 == (i,j) })){ //if used hashmap, it will update eg i = 1 again n again
 25                     heap.addPair((i,j,nums1[i]+nums2[j]))
 26                     
 27                     set.append((i,j))
 28                 }
 29                 
 30                 i = cur.0
 31                 j = cur.1 + 1
 32                 
 33                 if j < nums2.count && !(set.contains(where: { $0 == (i,j) })){
 34                     heap.addPair((i,j,nums1[i]+nums2[j]))
 35                     set.append((i,j))
 36                     
 37                 }
 38                 //dump(heap.pairs)
 39             }
 40             //print(set)
 41         }
 42        
 43         return result
 44     }
 45     
 46     struct PairHeap {
 47     
 48     var pairs: [(Int,Int,Int)] = []
 49     
 50     func getLeftChildIndex(_ parentIndex: Int) -> Int{
 51         return 2*parentIndex + 1
 52     }
 53     
 54     func getRightChildIndex(_ parentIndex: Int) -> Int{
 55         return 2*parentIndex + 2
 56     }
 57     
 58     func getParentIndex(_ childIndex: Int) -> Int{
 59         return (childIndex-1)/2
 60     }
 61     
 62     // Return Item From Heap
 63     func getLeftChild(_ parenIndex: Int) -> (Int,Int,Int) {
 64         return pairs[getLeftChildIndex(parenIndex)]
 65     }
 66     
 67     func getRightChild(_ parenIndex: Int) -> (Int,Int,Int) {
 68         return pairs[getRightChildIndex(parenIndex)]
 69     }
 70     
 71     func getParent(_ childIndex: Int) -> (Int,Int,Int) {
 72          return pairs[getParentIndex(childIndex)]
 73     }
 74     
 75     // Boolean Check
 76     private func hasLeftChild(_ index: Int) -> Bool {
 77         return getLeftChildIndex(index) < pairs.count
 78     }
 79     private func hasRightChild(_ index: Int) -> Bool {
 80         return getRightChildIndex(index) < pairs.count
 81     }
 82     private func hasParent(_ index: Int) -> Bool {
 83         return getParentIndex(index) >= 0
 84     }
 85     
 86     mutating func addPair(_ item : (Int,Int,Int)) {
 87         pairs.append(item)
 88         heapifyUp()
 89     }
 90     
 91     mutating public func poll() -> (Int,Int,Int) {
 92         if pairs.count > 0 {
 93             let item = pairs[0]
 94             pairs[0] = pairs[pairs.count - 1]
 95             pairs.removeLast()
 96             heapifyDown()
 97             return item
 98             
 99         } else {
100             fatalError()
101         }
102     }
103     
104     mutating func heapifyDown() {
105         var index = 0
106         
107         while hasLeftChild(index) {
108             var child = getLeftChildIndex(index)
109             if hasRightChild(index) && getRightChild(index).2 < pairs[child].2 {
110                 child = getRightChildIndex(index)
111             }
112             
113             if pairs[child].2 > pairs[index].2 {
114                 break
115             } else {
116                 pairs.swapAt(index, child)
117             }
118             index = child
119         }
120         
121     }
122     
123     mutating func heapifyUp() {
124         var index = pairs.count - 1
125         while hasParent(index) && getParent(index).2 > pairs[index].2 {
126             pairs.swapAt(index, getParentIndex(index))
127             index = getParentIndex(index)
128         }   
129     }    
130   }
131 }

---

884ms

  1 class Solution {
  2     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
  3         if nums1.count == 0 || nums2.count == 0 {
  4             return [[Int]]()
  5         }
  6         
  7         var rev = [[Int]]()
  8         var inHeap = Array(repeating: Array(repeating: false, count: nums2.count), count: nums1.count)
  9         var heap = Heap(type: .min)
 10         
 11         for i in 0 ..< nums2.count {
 12             let n = Node(pair: [nums1[0], nums2[i]], r: 0, c: i)
 13             heap.add(n)
 14         }
 15         
 16         while rev.count < k && heap.count > 0 {
 17             let n = heap.remove()!
 18             rev.append(n.pair)
 19             
 20             if n.r + 1 < nums1.count  {
 21                 heap.add(Node(pair: [nums1[n.r + 1], nums2[n.c]], r: n.r + 1, c: n.c))
 22             }
 23         }
 24         
 25         return rev
 26     }
 27 }
 28 
 29 struct Node: Comparable {
 30     let pair: [Int]
 31     let r: Int
 32     let c: Int
 33     
 34     static func == (lhs: Node, rhs: Node) -> Bool {
 35         return (lhs.pair[0] + lhs.pair[1]) == (rhs.pair[0] + rhs.pair[1])
 36     }
 37     
 38     static func < (lhs: Node, rhs: Node) -> Bool {
 39         return (lhs.pair[0] + lhs.pair[1]) < (rhs.pair[0] + rhs.pair[1])
 40     }
 41 }
 42 
 43 struct Heap {
 44     enum HeapType { case max, min }
 45 
 46     // an array representing the binary tree
 47     var tree = [E]()
 48 
 49     // indicating if this heap is a max heap or min heap
 50     let type: HeapType
 51 
 52     var count: Int {
 53         return tree.count
 54     }
 55 
 56     init(type: HeapType) {
 57         self.type = type
 58     }
 59 
 60     mutating func add(_ element: E) {
 61         tree.append(element)
 62         var child = tree.count - 1
 63         var parent = (child - 1) / 2
 64         while child > 0 && !satify(tree[parent], tree[child]) {
 65             tree.swapAt(parent, child)
 66             child = parent
 67             parent = (child - 1) / 2
 68         }
 69         assert(isHeap(0), "borken heap: \(tree)")
 70     }
 71 
 72     mutating func remove() -> E? {
 73         let rev = tree.first
 74         if tree.count > 0 {
 75             tree.swapAt(0, tree.count - 1)
 76             tree.removeLast()
 77             heapify(0)
 78         }
 79         assert(isHeap(0), "borken heap: \(tree)")
 80         return rev
 81     }
 82 
 83     func peek() -> E? {
 84         return tree.first
 85     }
 86 
 87     mutating private func heapify(_ rootIndex: Int) {
 88         let count = tree.count
 89         let leftChild = 2 * rootIndex + 1
 90         let rightChild = 2 * rootIndex + 2
 91 
 92         // no children
 93         if leftChild >= count { return }
 94 
 95         let targetChild = (rightChild >= count || satify(tree[leftChild], tree[rightChild])) ? leftChild : rightChild
 96         if (!satify(tree[rootIndex], tree[targetChild])) {
 97             tree.swapAt(rootIndex, targetChild)
 98             heapify(targetChild)
 99         }
100     }
101 
102     private func satify(_ lhs: E, _ rhs: E) -> Bool {
103         return (self.type == .min) ? (lhs <= rhs) : (lhs >= rhs)
104     }
105 
106     // debugging helper methods
107     private func isHeap(_ rootIndex: Int) -> Bool {
108         let leftChild = 2 * rootIndex + 1
109         let rightChild = 2 * rootIndex + 2
110         var valid = true
111         if leftChild < tree.count {
112             valid = satify(tree[rootIndex], tree[leftChild]) && isHeap(leftChild)
113         }
114 
115         if !valid { return false }
116 
117         if rightChild < tree.count {
118             valid = satify(tree[rootIndex], tree[rightChild]) && isHeap(rightChild)
119         }
120         return valid
121     }
122 }

---

1512ms

 1 class Solution {
 2     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
 3 
 4         var res = [[Int]]()
 5         for i in 0..<min(k, nums1.count) {
 6             for j in 0..<min(k, nums2.count) {
 7                 res.append([nums1[i], nums2[j]])
 8             }
 9         }
10         res.sort{$0[0]+$0[1] < $1[0]+$1[1]}
11         return Array(res.prefix(k))
12     }
13 }

---

2516ms

 1 class Solution {
 2     func kSmallestPairs(_ nums1: [Int], _ nums2: [Int], _ k: Int) -> [[Int]] {
 3         var res:[[Int]] = [[Int]]()
 4         let number1:Int = min(nums1.count, k)
 5         let number2:Int = min(nums2.count, k)
 6         for i in 0..<number1
 7         {
 8             for j in 0..<number2
 9             {
10                 res.append([nums1[i], nums2[j]])
11             }
12         }        
13         res.sort(by:{(a:[Int],b:[Int]) -> Bool in return a[0] + a[1] < b[0] + b[1] })        
14         if res.count > k
15         {
16             var ans:[[Int]] = [[Int]]()
17             for i in 0..<k
18             {                
19                 ans.append(res[i])
20             }
21             return ans
22         }
23         else
24         {
25             return res
26         }       
27     }
28 }

 

转载于:https://www.cnblogs.com/strengthen/p/10278047.html