列表

详情

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]

示例 3:

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

 

提示:

相似题目

有序矩阵中第 K 小的元素

找出第 K 小的数对距离

原站题解

去查看

上次编辑到这里，代码来自缓存 点击恢复默认模板
class Solution { public: vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) { } };

golang 解法, 执行用时: 152 ms, 内存消耗: 9.2 MB, 提交时间: 2022-11-22 10:15:08 

func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
    m, n := len(nums1), len(nums2)

    // 二分查找第 k 小的数对和
    left, right := nums1[0]+nums2[0], nums1[m-1]+nums2[n-1]+1
    pairSum := left + sort.Search(right-left, func(sum int) bool {
        sum += left
        cnt := 0
        i, j := 0, n-1
        for i < m && j >= 0 {
            if nums1[i]+nums2[j] > sum {
                j--
            } else {
                cnt += j + 1
                i++
            }
        }
        return cnt >= k
    })

    // 找数对和小于 pairSum 的数对
    i := n - 1
    for _, num1 := range nums1 {
        for i >= 0 && num1+nums2[i] >= pairSum {
            i--
        }
        for _, num2 := range nums2[:i+1] {
            ans = append(ans, []int{num1, num2})
            if len(ans) == k {
                return
            }
        }
    }

    // 找数对和等于 pairSum 的数对
    i = n - 1
    for _, num1 := range nums1 {
        for i >= 0 && num1+nums2[i] > pairSum {
            i--
        }
        for j := i; j >= 0 && num1+nums2[j] == pairSum; j-- {
            ans = append(ans, []int{num1, nums2[j]})
            if len(ans) == k {
                return
            }
        }
    }
    return
}

python3 解法, 执行用时: 1036 ms, 内存消耗: 29.9 MB, 提交时间: 2022-11-22 10:14:16 

'''
二分查找
'''
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)

        # 二分查找第 k 小的数对和
        left, right = nums1[0] + nums2[0], nums1[m - 1] + nums2[n - 1] + 1
        while left < right:
            mid = (left + right) // 2
            cnt = 0
            i, j = 0, n - 1
            while i < m and j >= 0:
                if nums1[i] + nums2[j] > mid:
                    j -= 1
                else:
                    cnt += j + 1
                    i += 1
            if cnt < k:
                left = mid + 1
            else:
                right = mid
        pairSum = left

        ans = []
        # 找数对和小于 pairSum 的数对
        i = n - 1
        for num1 in nums1:
            while i >= 0 and num1 + nums2[i] >= pairSum:
                i -= 1
            for j in range(i + 1):
                ans.append([num1, nums2[j]])
                if len(ans) == k:
                    return ans

        # 找数对和等于 pairSum 的数对
        i = n - 1
        for num1 in nums1:
            while i >= 0 and num1 + nums2[i] > pairSum:
                i -= 1
            j = i
            while j >= 0 and num1 + nums2[j] == pairSum:
                ans.append([num1, nums2[j]])
                if len(ans) == k:
                    return ans
                j -= 1
        return ans

python3 解法, 执行用时: 196 ms, 内存消耗: 30 MB, 提交时间: 2022-11-22 10:13:14 

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        pq = [(nums1[i] + nums2[0], i, 0) for i in range(min(m, k))]
        ans = []
        while pq and len(ans) < k:
            _, i, j = heapq.heappop(pq)
            ans.append([nums1[i], nums2[j]])
            if j+1 < n:
                heapq.heappush(pq, (nums1[i]+nums2[j+1], i, j+1))
        return ans

golang 解法, 执行用时: 160 ms, 内存消耗: 9.8 MB, 提交时间: 2022-11-22 10:08:20 

func kSmallestPairs(nums1, nums2 []int, k int) (ans [][]int) {
    m, n := len(nums1), len(nums2)
    h := hp{nil, nums1, nums2}
    for i := 0; i < k && i < m; i++ {
        h.data = append(h.data, pair{i, 0})
    }
    for h.Len() > 0 && len(ans) < k {
        p := heap.Pop(&h).(pair)
        i, j := p.i, p.j
        ans = append(ans, []int{nums1[i], nums2[j]})
        if j+1 < n {
            heap.Push(&h, pair{i, j + 1})
        }
    }
    return
}

type pair struct{ i, j int }
type hp struct {
    data         []pair
    nums1, nums2 []int
}
func (h hp) Len() int            { return len(h.data) }
func (h hp) Less(i, j int) bool  { a, b := h.data[i], h.data[j]; return h.nums1[a.i]+h.nums2[a.j] < h.nums1[b.i]+h.nums2[b.j] }
func (h hp) Swap(i, j int)       { h.data[i], h.data[j] = h.data[j], h.data[i] }
func (h *hp) Push(v interface{}) { h.data = append(h.data, v.(pair)) }
func (h *hp) Pop() interface{}   { a := h.data; v := a[len(a)-1]; h.data = a[:len(a)-1]; return v }

上一题

© 2024 leetcode.cn