Appearance
373. 查找和最小的K对数字
题目描述
给定两个以 非递减顺序排列 的整数数组 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]
提示:
1 <= nums1.length, nums2.length <= 105
-109 <= nums1[i], nums2[i] <= 109
- `nums1` 和 `nums2` 均为 **升序排列**
1 <= k <= 104
- `k <= nums1.length * nums2.length`
难度： Medium
题解代码
javascript
/**
* @param {number[]} nums1
* @param {number[]} nums2
* @param {number} k
* @return {number[][]}
*/
// 暴力堆--超时
var kSmallestPairs = function(nums1, nums2, k) {
const heap = []
let count = 0
for (let i = 0; i < nums1.length; i++) {
for (let j = 0; j < nums2.length; j++) {
if (count < k) {
heap.push([nums1[i], nums2[j]])
if (count === k - 1) genHeap(heap, k, 0)
} else {
if ((heap[0][0] + heap[0][1]) > (nums1[i] + nums2[j])) {
heap[0] = [nums1[i], nums2[j]]
maxHeapify(heap, k, 0)
}
}
count++
}
}
return heap
function genHeap(arr, size) {
for (let i = Math.floor(arr.length / 2); i >= 0; i--) {
maxHeapify(arr, size, i)
}
}
function maxHeapify(arr, size, i) {
const l = i * 2 + 1
const r = i * 2 + 2
let max = i
const val = idx => arr[idx][0] + arr[idx][1]
if (l < size && val(l) > val(max)) {
max = l
}
if (r < size && val(r) > val(max)) {
max = r
}
if (max !== i) {
[arr[max], arr[i]] = [arr[i], arr[max]]
maxHeapify(arr, size, max)
}
}
};
// 手写堆
var kSmallestPairs = function(nums1, nums2, k) {
const heap = []
const res = []
const len = Math.min(nums1.length, k)
for (let i = 0; i < len; i++) {
heap.push([i, 0])
}
genHeap(heap)
while ( res.length < k && heap.length) {
const [i, j] = heap.shift()
res.push([nums1[i], nums2[j]])
if (j + 1 < nums2.length) heap.unshift([i, j + 1])
minHeapify(heap, 0)
}
return res
function genHeap (heap) {
for (let i = Math.floor(heap.length / 2); i >= 0; i--) {
minHeapify(heap, i)
}
}
function minHeapify (heap, i) {
const l = i * 2 + 1
const r = i * 2 + 2
let min = i
const size = heap.length
const val = idx => nums1[heap[idx][0]] + nums2[heap[idx][1]]
if (l < size && val(l) < val(min)) {
min = l
}
if (r < size && val(r) < val(min)) {
min = r
}
if (min !== i) {
[heap[i], heap[min]] = [heap[min], heap[i]]
minHeapify(heap, min)
}
}
};
// 内置堆
var kSmallestPairs = function (nums1, nums2, k) {
const res = [];
// leetcode 内置API
const pq = new MinPriorityQueue({
compare: (a, b) => nums1[a[0]] + nums2[a[1]] - (nums1[b[0]] + nums2[b[1]])
});
for (let i = 0; i < Math.min(k, nums1.length); i++) pq.enqueue([i, 0]);
while (res.length < k && pq.size()) {
const [i, j] = pq.dequeue();
if (j + 1 < nums2.length) pq.enqueue([i, j + 1]);
res.push([nums1[i], nums2[j]]);
}
return res;
};