【Lintcode】465. Kth Smallest Sum In Two Sorted Arrays

题目地址：

https://www.lintcode.com/problem/kth-smallest-sum-in-two-sorted-arrays/description

给定两个递增数组$A$和$B$，从两个数组各自挑一个数求和，问这些和里排名第$k$的是多少（大小一样的话，有多少个就占多少个排名。比如和的可能取值是$(0,1,1,2,...)$，那么排名第$3$的和就是$1$）。

可以将$A$和$B$张成一个和矩阵，具体来说，如果$A=[a_0,a_1,...,a_{n-1}]$，$B=[b_0,b_1,...,b_{m-1}]$，那么它们张成的矩阵就是：$$\left[\begin{array}{cccc}{a}_0+b_0& a_1+b_0& ...& a_{n-1}+b_0\\ a_0+b_1& a_1+b_1& ...& a_{n-1}+b_1\\ a_0+b_2& a_1+b_2& ...& a_{n-1}+b_2\\ ...& ...& ...& ...\\ a_0+b_{m-1}& a_1+b_{m-1}& ...& a_{n-1}+b_{m-1}\end{array}\right]$$接下来，就可以用https://blog.csdn.net/qq_46105170/article/details/105592678里面的方法来做了。

法1：二分答案。很显然答案的范围就是$[A[0]+B[0],A[n-1]+B[n-1]]$，我们需要求出某个数$x$，它是使得小于等于$x$的和的个数大于等于$k$的最小数。显然这样的$x$就是所求的答案。而小于等于$x$的和的个数，可以参考给的连接的那道题。代码如下：

public class Solution {
    /**
     * @param A: an integer arrays sorted in ascending order
     * @param B: an integer arrays sorted in ascending order
     * @param k: An integer
     * @return: An integer
     */
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // write your code here
        int l = A[0] + B[0], r = A[A.length - 1] + B[B.length - 1];
        while (l < r) {
            int m = l + (r - l >> 1);
            if (countLessEq(A, B, m) >= k) {
                r = m;
            } else {
                l = m + 1;
            }
        }
        
        return l;
    }
    
    private int countLessEq(int[] A, int[] B, int target) {
        int count = 0;
        int x = B.length - 1, y = 0;
        while (x >= 0 && y < A.length) {
            if (A[y] + B[x] <= target) {
                count += x + 1;
                y++;
            } else {
                x--;
            }
        }
        
        return count;
    }
}

时间复杂度$O((n+m)\log (A[n-1]+B[m-1]-A[0]-B[0]))$，空间$O(1)$。

法2：多路归并。思路参考上面的链接。代码如下：

import java.util.PriorityQueue;

public class Solution {
    
    class Pair {
        int x, y, val;
    
        public Pair(int x, int y, int val) {
            this.x = x;
            this.y = y;
            this.val = val;
        }
    }
    
    /**
     * @param A: an integer arrays sorted in ascending order
     * @param B: an integer arrays sorted in ascending order
     * @param k: An integer
     * @return: An integer
     */
    public int kthSmallestSum(int[] A, int[] B, int k) {
        // write your code here
        PriorityQueue<Pair> minHeap = new PriorityQueue<>((p1, p2) -> Integer.compare(p1.val, p2.val));
        for (int i = 0; i < B.length; i++) {
            minHeap.offer(new Pair(i, 0, B[i] + A[0]));
        }
    
        for (int i = 1; i <= k - 1; i++) {
            Pair cur = minHeap.poll();
            if (cur.y < A.length - 1) {
                minHeap.offer(new Pair(cur.x, cur.y + 1, B[cur.x] + A[cur.y + 1]));
            }
        }
        
        return minHeap.peek().val;
    }
}

时间复杂度$O(k\log m)$，空间$O(m)$。

注解：由复杂度可以知道，我们要把较短的数组视为$B$，复杂度会比较低。