Kth Smallest Sum In Two Sorted Arrays

Posted 2020-07-28 Sheryl Wang

Given two integer arrays sorted in ascending order and an integer k. Define sum = a + b, where a is an element from the first array and b is an element from the second one. Find the kth smallest sum out of all possible sums.

Given [1, 7, 11] and [2, 4, 6].

For k = 3, return 7.

For k = 4, return 9.

For k = 8, return 15.

这题是Kth Smallest Number in Sorted Matrix的follow up.一开始可能很难想到，但是这两个数组的和如果按照第一个数组的第一个数组与第二个数组的所有元素配对，第二个元素与第二个数组的所有元素配对．则形成的就是一个排序矩阵．所以完全可以采取之前的heap思路来做．但是这题如果直接采用之前的思路来做会出现超内存,我们反思一下之前的做法. 对于矩阵来说,长宽不会特别大.但是对于数组,长度可以非常长,甚至远远大于k的大小.所以在这时候,完全没有必要处理到矩阵的k行和k列以外去.visited矩阵的大小可以限定为min(k,len(matrix))和min(k,len(matrix[0])). 大大减小了空间.空间复杂度为O(min(k,m)*min(k,n)).时间复杂度为O(k*log(m,n,k)).代码如下:

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):
        
        a = len(A)
        b = len(B)
        from heapq import *
        heap = [(A[0]+B[0],(0,0))]
        visited = [[False] * min(b,k) for j in xrange(min(a,k))]
        i = 0
        while i < k:
            cur, (x,y) = heappop(heap)
            i += 1
            if x + 1 < min(a, k) and visited[x+1][y] == False:
                heappush(heap,(A[x+1] + B[y], (x+1,y)))
                visited[x+1][y] = True
            if y + 1 < min(b, k) and visited[x][y+1] == False:
                heappush(heap, (A[x] + B[y+1], (x, y+1)))
                visited[x][y+1] = True
                
        return cur

另外一种处理方法是我们先在堆中放入第一行的元素,之后每次朝下更新,这样避免了向右和向下带来的重复,不需要维护visited矩阵和进行繁复的边界判断.代码参考九章如下:

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
        import heapq
        m, n = len(A), len(B)
        def vertical_search(k):
            heap = []
            for i in range(min(k, n)):
                heapq.heappush(heap, (A[0] + B[i], 0, i))
            while k > 1:
                min_element = heapq.heappop(heap)
                x, y = min_element[1], min_element[2]
                if x + 1 < m:
                    heapq.heappush(heap, (A[x + 1] + B[y], x + 1, y))
                k -= 1
            return heapq.heappop(heap)[0]
        return vertical_search(k)