描述

给定两个排好序的数组 A, B，定义集合 sum = a + b ，求 sum 中第k小的元素

样例

给出 A = [1,7,11] B = [2,4,6]
当 k = 3, 返回 7.
当 k = 4, 返回 9.
当 k = 8, 返回 15.

挑战

Do it in either of the following time complexity:

1. O(k log min(n, m, k)). where n is the size of A, and m is the size of B.
2. O( (m + n) log maxValue). where maxValue is the max number in A and B.

思路

排序矩阵中的从小到大第k个数这题的马甲变换

代码

1. 两个排序数组分别作为行列，对应每个位置的和组成的矩阵其实是按照行从小到大列从小到大的顺序排列的

class Pair {
    public int x, y, sum;
    public Pair(int x, int y, int val) {
        this.x = x;
        this.y = y;
        this.sum = val;
    }
}
class PairComparator implements Comparator<Pair> {
    public int compare(Pair a, Pair b) {
        return a.sum - b.sum;
    }
}

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) {
        int[] dx = new int[]{0, 1};
        int[] dy = new int[]{1, 0};
        boolean[][] hash = new boolean[A.length][B.length];
        PriorityQueue<Pair> minHeap = new PriorityQueue<Pair>(k, new PairComparator());
        minHeap.add(new Pair(0, 0, A[0] + B[0]));

        for(int i = 0; i < k - 1; i ++){
            Pair cur = minHeap.poll();
            for(int j = 0; j < 2; j ++){
                int next_x = cur.x + dx[j];
                int next_y = cur.y + dy[j];
                Pair next_Pair = new Pair(next_x, next_y, 0);
                if(next_x < A.length && next_y < B.length &&  !hash[next_x][next_y]){
                    hash[next_x][next_y] = true;
                    next_Pair.sum = A[next_x] + B[next_y];
                    minHeap.add(next_Pair);
                }
            }
        }
        return minHeap.peek().sum;
    }
}

2. 二分法(看不太懂)
   复杂度: O((m + n) * lg (MAX_SUM - MIN_SUM))
   思路:二分法相当于元素到其排序的映射，这里的元素是两个数组任意两个元素的和

class Range {
    int lt = 0;
    int le = 0;
}
    

    Range search (int[] A, int[] B, int v) {
        int i = A.length - 1;
        int j = 0;
        Range ans = new Range();
        
        //  find number elements that are smaller than v
        while (i >= 0 && j < B.length) {
            int sum = A[i] + B[j];
            if (sum < v) {
                j ++;
                ans.lt += i + 1;
            } else {
                i--;
            }
        }
        
        i = A.length - 1;
        j = 0;
        
        //  find number of elements that are less than or equal to v
        while (i >= 0 && j < B.length) {
            int sum = A[i] + B[j];
            if (sum <= v) {
                j ++;
                ans.le += i + 1;
            } else {
                i --;
            }
        }
        return ans;
    }


    public int kthSmallestSum (int[] A, int[] B, int k) {
        
        int min = A[0] + B[0];
        Range ret = search(A, B, min);
        if (ret.le >= k) return min;
        
        int max = A[A.length - 1] + B[B.length - 1];
        ret = search(A, B, max);
        if (ret.lt < k) return max;
        
        while (min < max - 1) {  // invariant: min < X < max.
            
            int mid = min + (max - min) / 2;
            ret = search(A, B, mid);
            
            if (ret.lt < k && ret.le >= k) return mid;
            if (ret.le < k) min = mid;
            if (ret.le >= k) max = mid;
        }
        
        throw new IllegalArgumentException("k = " + k);
    }