LintCode 465 [Kth Smallest Sum in Two Sorted Array]

原题

给定两个排好序的数组 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]