堆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大的思路！！！}$

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))