LeetCode in Kotlin
373. Find K Pairs with Smallest Sums
Medium
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.
Return 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]
Constraints:
1 <= nums1.length, nums2.length <= 105-109 <= nums1[i], nums2[i] <= 109nums1andnums2both are sorted in ascending order.1 <= k <= 104
Solution
import java.util.PriorityQueue
class Solution {
private class Node(index: Int, num1: Int, num2: Int) {
var sum: Long
var al: MutableList<Int>
var index: Int
init {
sum = num1.toLong() + num2.toLong()
al = ArrayList()
al.add(num1)
al.add(num2)
this.index = index
}
}
fun kSmallestPairs(nums1: IntArray, nums2: IntArray, k: Int): List<List<Int>> {
val queue = PriorityQueue { a: Node, b: Node -> if (a.sum < b.sum) -1 else 1 }
val res: MutableList<List<Int>> = ArrayList()
run {
var i = 0
while (i < nums1.size && i < k) {
queue.add(Node(0, nums1[i], nums2[0]))
i++
}
}
var i = 1
while (i <= k && queue.isNotEmpty()) {
val cur = queue.poll()
res.add(cur.al)
val next = cur.index
val lastNum1 = cur.al[0]
if (next + 1 < nums2.size) {
queue.add(Node(next + 1, lastNum1, nums2[next + 1]))
}
i++
}
return res
}
}