原题
给定两个排好序的数组 A, B，定义集合 sum = a + 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.
当 k = 8, 返回 15.
解题思路
- 使用set来记录访问过的pair
visited = set()
visited.add(10)
10 is in visited # return True
- 使用一个minHeap，每次可以取最小值。放入堆中的数据结构为(sum, i, j)
- 每次从堆中取出一个最小值，下一轮的candidates将是(sum, i+1, j), (sum, i, j+1)，将这两个candidates加入的heap之间要去重，因为有可能重复
- 从堆中取k次之后，就拿到了想要的结果
完整代码
import Queue
class Solution:
# @param {int[]} A an integer arrays sorted in ascending order
# @param {int[]} B an integer arrays sorted in ascending order
# @param {int} k an integer
# @return {int} an integer
def kthSmallestSum(self, A, B, k):
# Write your code here
if not A or not B:
return 0
myQueue = Queue.PriorityQueue()
visited = set()
myQueue.put((A[0]+B[0], 0, 0))
visited.add((0, 0))
while k > 1:
cur = myQueue.get()
k -= 1
if cur[1] + 1 < len(A) and (cur[1] + 1, cur[2]) not in visited:
visited.add((cur[1] + 1, cur[2]))
myQueue.put((A[cur[1] + 1] + B[cur[2]], cur[1] + 1, cur[2]))
if cur[2] + 1 < len(B) and (cur[1], cur[2] + 1) not in visited:
visited.add((cur[1], cur[2] + 1))
myQueue.put((A[cur[1]] + B[cur[2] + 1], cur[1], cur[2] + 1))
res = myQueue.get()
return res[0]
