本帖已被设为精华帖！
- 问题描述： 给定两个排好序的数组 A，B，定义集合 sum = a + b, 其中 a 来自数组 A，b 来自数组 B，求 sum 中第 k 小的元素
- 问题示例： A = [1, 7, 11] , B = [2, 4, 6] sum = [3, 5, 7, 9, 11, 13, 13, 15, 17] 当 k = 3，返回 7 当 k = 4, 返回 9
参考代码:
class Solution:
def kthSmallestSum(self, A, B, k):
if not A or not B:
return None
n, m = len(A), len(B)
minheap = [A[0] + B[0], 0, 0]
visited = set([0])
num = None
for _ in range(k):
num, x, y = heapq.heappop(minheap)
if x+1 < n and (x + 1) * m + y not in visited:
heapq.heappush(minheap, (A[x + 1] + B[y], x + 1, y))
visited.add((x + 1) * m + y)
if y+1 < m and x * m + y + 1 not in visited:
heapq.heappush(minheap, (A[x] + B[y + 1], x, y+1))
visited.add(x * m + y + 1)
return num