- 两个排序数组和的第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.挑战
挑战一下更小的时间复杂度做法：
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))
