Hao.Jia
文章
59
分类
15
评论
0

每日一题：373. 查找和最小的 K 对数字（堆）

| 算法 • 题解 •
字数总计 8936 | 阅读时长 22分钟 | 阅读量 224

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]

提示:

1 <= nums1.length, nums2.length <= 1e5
-109 <= nums1[i], nums2[i] <= 1e9
nums1 和 nums2 均为升序排列
1 <= k <= 1000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

题解:

方法一（k2logk）：

已知nums1和nums2递增且有序，则前k小的组合一定出现在nums1[1-k]和nums2[1-k]的组合中，组合的个数为k2,然后把这k2的所有情况枚举出来，维护一个k大小的堆，这也是topk问题的解决思路。
证明：为什么一定前k小的组合一定是由nums1[1-k] 组合 nums2[1-k] 而成？

  • 取nums1[k+1],nums2[k+1]，显然存在[nums1[i]+nums2[i]]<[nums1[k+1]+nums2[k+1]] (i<k+1)，i的取值有k这种情况，所以nums1[k+1],nums2[k+1]比不可能是前k小的组合
  • 取nums1[i],nums2[j]（i<k,j>k）显然存在[nums1[i]+nums2[q]]<[nums1[i]+nums2[j]] (q<j),q的取值有j-1个取值，故nums1[i],nums2[j]（i<k,j>k）必不可能是前k小的组合
  • nums1[i],nums2[j]（i>k,j<k）情况与上相同，从而得证

代码：
自写小根堆

class Solution:

    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap=Heap(k)
        for i in nums1[:k]:
            for j in nums2[:k]:
                heap.add((i,j,i+j))
        res=[]
        for i in range(k):
            max_ = heap.pop()
            if max_== None:
                break
            res.append([max_[0],max_[1]])
        res.reverse()
        return res

class Heap:
    def __init__(self, k):
        self.heap = [(0, 0, 0)] * (k + 1)
        self.size = 0
        self.max_size = k

    def pop(self):
        if self.size==0:
            return
        tmp = self.heap[self.size]
        self.heap[self.size] = self.heap[1]
        self.size -= 1
        now = 1
        while now * 2 <= self.size:
            a = now * 2
            b = now * 2 + 1
            if b > self.size:
                if self.heap[a][2] > tmp[2]:
                    self.heap[now] = self.heap[a]
                    now = a
                else:
                    break
            else:
                if self.heap[a][2] < self.heap[b][2]:
                    a = b
                if self.heap[a][2] > tmp[2]:
                    self.heap[now] = self.heap[a]
                    now = a
                else:
                    break
        self.heap[now] = tmp
        return self.heap[self.size + 1]

    def add(self, node):
        if self.size == self.max_size :
            if node[2] < self.top()[2]:
                self.pop()
            else:
                return
        self.size += 1
        self.heap[self.size] = node
        self.adjust()

    def adjust(self):
        now = self.size
        tmp = self.heap[now]
        while now > 1:
            if tmp[2] > self.heap[now // 2][2]:
                self.heap[now] = self.heap[now // 2]
                now = now // 2
            else:
                break
        self.heap[now] = tmp

    def top(self):
        return self.heap[1]

image.png
python自带堆heapq操作

import heapq
class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap=[]
        for i in nums1[:k]:
            for j in nums2[:k]:
                heapq.heappush(heap,(i+j,[i,j]))
        res=[]
        for i in range(min(k,len(heap))):

            max_ = heapq.heappop(heap)
            if max_== None:
                break
            res.append(max_[1])
        return res

image.png

方法二（klogk）：

两个数组都是升序序列 ，所以，我们可以利用这个条件提高效率。

我们可以使用优先级队列来实现，队列中保存的是 [index1, index2]，分别表示 nums1 的索引和 nums2 的索引，初始时把 [0,0]、[1, 0]、[2, 0]、…… 入队，即让 nums2 的索引全部从 0 开始，每次弹出 nums1[index1] + nums2[index2] 较小者。弹出之后，再把 index2 后移一位继续添加到优先级队列中，依次往复，最终弹出 k 次就是我们的结果。

为什么这样是可行的？

首先，利用两个数组都是升序的条件，结果中最小的组合肯定是 nums1[0] + nums2[0]。

但是，次小的是什么呢？ 是 nums1[0] + nums2[1] 还是 nums1[1] + nums2[0] 呢？我们不知道。

所以，我们需要设计一种比较方式使我们能知道上述比较的结果，使用优先级队列就可以解决。

假设，我们以 [0, 1] 表示 nums1[0] + nums2[1] 的结果，整个过程的处理如下：

先把所有的 nums1 的索引入队，即入队的元素有 [0, 0]、[1, 0]、[2, 0]、[3, 0]、…
首次弹出的肯定是 [0, 0]，再把 [0, 1] 入队；
这样就可以通过优先级队列比较 [0, 1] 和 [1, 0] 的结果，再弹出较小者；
依次进行，进行 k 轮。

作者：tong-zhu
链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/solution/tong-ge-lai-shua-ti-la-you-xian-ji-dui-l-fw7y/
来源：力扣（LeetCode）
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap = Heap(k * 2)
        i = 0
        for j in range(min(k, len(nums2))):
            heap.add((i, j, nums1[i] + nums2[j]))
        res = []
        for i in range(k):
            # print(heap.heap)
            min_ = heap.pop()
            if min_==None:
                break
            res.append([nums1[min_[0]], nums2[min_[1]]])
            if min_[0] + 1 >= len(nums1) and min_[1] + 1 >= len(nums2):
                continue
            elif min_[0] + 1 >= len(nums1):
                heap.add((min_[0], min_[1] + 1, nums1[min_[0]] + nums2[min_[1] + 1]))
            elif min_[1] + 1 >= len(nums2):
                heap.add((min_[0] + 1, min_[1], nums1[min_[0] + 1] + nums2[min_[1]]))
            else:
                heap.add((min_[0] + 1, min_[1], nums1[min_[0] + 1] + nums2[min_[1]]))
                heap.add((min_[0], min_[1] + 1, nums1[min_[0]] + nums2[min_[1] + 1]))

        return res


class Heap:
    def __init__(self, k):
        self.heap = [(0, 0, 0)] * (k + 1)
        self.size = 0
        self.max_size = k
        self.visited = set([])

    def pop(self):
        if self.size == 0:
            return
        tmp = self.heap[self.size]
        self.heap[self.size] = self.heap[1]
        self.size -= 1
        now = 1
        while now * 2 <= self.size:
            a = now * 2
            b = now * 2 + 1
            if b > self.size:
                if self.heap[a][2] < tmp[2]:
                    self.heap[now] = self.heap[a]
                    now = a
                else:
                    break
            else:
                if self.heap[a][2] > self.heap[b][2]:
                    a = b
                if self.heap[a][2] < tmp[2]:
                    self.heap[now] = self.heap[a]
                    now = a
                else:
                    break
        self.heap[now] = tmp
        return self.heap[self.size + 1]

    def add(self, node):
        # if self.size == self.max_size :
        #     self.pop()
        if node in self.visited:
            return
        self.visited.add(node)
        self.size += 1
        self.heap[self.size] = node
        self.adjust()

    def adjust(self):
        now = self.size
        tmp = self.heap[now]
        while now > 1:
            if tmp[2] < self.heap[now // 2][2]:
                self.heap[now] = self.heap[now // 2]
                now = now // 2
            else:
                break
        self.heap[now] = tmp

    def top(self):
        return self.heap[1]

image.png

文章作者： Hao.Jia
本文链接：
版权声明： 本站所有文章除特别声明外，均采用 CC BY-NC-SA 4.0 许可协议。转载请注明来自 Hao.Jia's Blog！
算法 题解
喜欢就支持一下吧