原題
給定兩個(gè)排好序的數(shù)組 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]
當(dāng) k = 3, 返回 7.
當(dāng) k = 4, 返回 9.
當(dāng) k = 8, 返回 15.
解題思路
- 使用set來(lái)記錄訪問(wèn)過(guò)的pair
visited = set()
visited.add(10)
10 is in visited # return True
- 使用一個(gè)minHeap周崭，每次可以取最小值哑诊。放入堆中的數(shù)據(jù)結(jié)構(gòu)為(sum, i, j)
- 每次從堆中取出一個(gè)最小值撒璧，下一輪的candidates將是(sum, i+1, j), (sum, i, j+1)透葛，將這兩個(gè)candidates加入的heap之間要去重，因?yàn)橛锌赡苤貜?fù)
- 從堆中取k次之后卿樱，就拿到了想要的結(jié)果
完整代碼
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]
