堆heap的使用——两个排序数组和的第K小

2. 两个排序数组和的第K小 · Kth Smallest Sum In Two Sorted Arrays
   堆/优先队列Sorted Matrix优先队列堆

   描述

   给定两个排好序的数组 A, B，定义集合 sum = a + b ，其中a来自A数组，b来自B数组，求 sum 中第k小的元素

   样例

   样例1

   输入:
   a = [1, 7, 11]
   b = [2, 4, 6]
   k = 3
   输出: 7
   说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第三个是7.
   

   样例2

   输入:
   a = [1, 7, 11]
   b = [2, 4, 6]
   k = 4
   输出: 9
   说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第四个是9.
   

   样例3

   输入:
   a = [1, 7, 11]
   b = [2, 4, 6]
   k = 8
   输出: 15
   说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第八个是15.
   

   挑战

   挑战一下更小的时间复杂度做法：

   1. $\mathrm{O(klogmin(n,m,k))}$
   O(klogmin(n,m,k)​), $\mathrm{nn\; 是A\; 数组的大小，m}$
3. m是B 数组的大小
4. $\mathrm{O((m+n)logmaxValue)}$

O((m+n)logmaxValue​),$\mathrm{maxValuemaxValue是\; A\; 数组和\; B\; 数组中的最大值}$

 

$\mathrm{使用的是递增矩阵中第K大的思路！！！}$

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

import collections import heapq   class Solution:     """     输入: a = [1, 7, 11] b = [2, 4, 6] ==> a+b, 1+2, 1+4, 1+6, 7+2, k = 3 输出: 7 说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第三个是7.   样例2   输入: a = [1, 7, 11] b = [2, 4, 6] k = 4 输出: 9 说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第四个是9.   输入: a = [1, 7, 11] b = [2, 4, 6, 9] 1+2=3, 1+4=5, 1+6=7, 1+9=10 7+2=9   1+2, 1+4<7+2(5),1+6<7+2(7),1+9>7+2(9),1+9<7+4(10),7+4, repeat k = 8 输出: 15 说明: 满足条件的所有的和有[3, 5, 7, 9, 11, 13, 13, 15, 17]，其中第八个是15.     """     def findK(self, a, b, k):         seen = {(0, 0)}         ans = -1         q = [(a[0]+b[0], 0, 0)]         for n in range(k):             item, i, j = heapq.heappop(q)             ans = item             if i+1 < len(a):                 if (i+1, j) not in seen:                     heapq.heappush(q, (a[i+1]+b[j], i+1, j))                     seen.add((i+1, j))             if j+1 < len(b):                 if (i, j+1) not in seen:                     heapq.heappush(q, (a[i]+b[j+1], i, j+1))                     seen.add((i, j+1))           return ans   a = [1, 7, 11] b = [2, 4, 6, 9]   c = [] for i in a:     for j in b:         c.append(i+j) print(sorted(c))   print(Solution().findK(a, b, k=3)) print(Solution().findK(a, b, k=4)) print(Solution().findK(a, b, k=8)) print(Solution().findK(a, b, k=11)) print(Solution().findK(a, b, k=12))