描述

給定兩個(gè)排好序的數(shù)組 A, B险绘，定義集合 sum = a + b 商模，求 sum 中第k小的元素

樣例

給出 A = [1,7,11] B = [2,4,6]
當(dāng) k = 3, 返回 7.
當(dāng) k = 4, 返回 9.
當(dāng) k = 8, 返回 15.

挑戰(zhàn)

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個(gè)數(shù)這題的馬甲變換

代碼

1. 兩個(gè)排序數(shù)組分別作為行列，對(duì)應(yīng)每個(gè)位置的和組成的矩陣其實(shí)是按照行從小到大列從小到大的順序排列的

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. 二分法(看不太懂)
   復(fù)雜度: O((m + n) * lg (MAX_SUM - MIN_SUM))
   思路:二分法相當(dāng)于元素到其排序的映射厦幅，這里的元素是兩個(gè)數(shù)組任意兩個(gè)元素的和

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);
    }