LintCode 465. Kth Smallest Sum In Two Sorted Arrays (LintCode 401变种，堆经典题)

465. Kth Smallest Sum In Two Sorted Arrays
     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.

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

For k = 3, return 7.

For k = 4, return 9.

For k = 8, return 15.

Challenge
Do it in either of the following time complexity:

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

代码如下:
这题我尝试用双指针好像比较麻烦，因为指针有时要往后移动。
以A = [1,7,11], B = [2,4,6]为例。先pa = 0 (即A[0]=1) 不动，pB=0, 1, 2，故sum = 3,5, 7。但然后pA移动到1, pB还得移回到0。时间复杂度高。

解法1：我们把A[i]+B[j]构成一个矩阵，这样就变成
LintCode 401. Kth Smallest Number in Sorted Matrix 这题了。利用行列已经有序这个特性，用最小堆。
时间复杂度O(KlogK) ，严格说是O(KlogMin(n,m,k))。

struct Node {
    int row;
    int col;
    int val;
    Node (int r, int c, int v) : row(r), col(c), val(v) {}
};

struct cmp {
    bool operator() (const Node & a, const Node & b) {
        return a.val >= b.val;   //minHeap
    }
};

class Solution {
public:
    /**
     * @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
     */
    int kthSmallestSum(vector<int> &A, vector<int> &B, int k) {
        int lenA = A.size();
        int lenB = B.size();
        if (lenA == 0 || lenB == 0) return 0;

        vector<vector<int>> visited(lenA, vector<int>(lenB, 0));
        priority_queue<Node, vector<Node>, cmp> minHeap; 
        minHeap.push(Node(0, 0, A[0] + B[0]));
        
        for (int i = 1; i < k; ++i) {
            Node topNode = minHeap.top();
            minHeap.pop();  //totally popped k - 1 elems

            int newR = topNode.row + 1;
            if (newR < lenA && !visited[newR][topNode.col]) {
                minHeap.push(Node(newR, topNode.col, A[newR] + B[topNode.col]));
                visited[newR][topNode.col] = 1;
            }
            
            int newC = topNode.col + 1;
            if (newC < lenB && !visited[topNode.row][newC]) {
                minHeap.push(Node(topNode.row, newC, A[topNode.row] + B[newC]));
                visited[topNode.row][newC] = 1;
            }
        }
        
        return minHeap.top().val;
    }
};

二刷：

struct Node {
    int x;
    int y;
    int v;
    Node(int row, int col, int val) : x(row), y(col), v(val) {}
    bool operator< (const Node &node) const {
        return v > node.v;    //minHeap;
    }
};

class Solution {
public:
    /**
     * @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
     */
    int kthSmallestSum(vector<int> &a, vector<int> &b, int k) {
        int aSize = a.size();
        int bSize = b.size();
        priority_queue<Node> minHeap;
        vector<vector<int>> visited(aSize, vector<int>(bSize, false));
        minHeap.push(Node(0, 0, a[0] + b[0]));

        for (int i = 1; i < k; i++) {
            Node topNode = minHeap.top();
            minHeap.pop();
 
            if (topNode.x + 1 < aSize && !visited[topNode.x + 1][topNode.y]) {
                minHeap.push(Node(topNode.x + 1, topNode.y, a[topNode.x + 1] + b[topNode.y]));
                visited[topNode.x + 1][topNode.y] = true;
            }

            if (topNode.y + 1 < bSize && !visited[topNode.x][topNode.y + 1]) {
                minHeap.push(Node(topNode.x, topNode.y + 1, a[topNode.x] + b[topNode.y + 1]));
                visited[topNode.x][topNode.y + 1] = true;
            }            
        }   
        
        return minHeap.top().v;
    }
};

解法2：用大顶堆，但是加上b数组已经有序这个特点。注意我们在添加a[i]+b[j]的时候，如果maxHeap已经有k个元素，并且a[i]+b[j]比maxHeap.top()还大，那么就没有必要push {a[i], b[j]}，并且b数组后面的b[j+1],b[j+2]…也不用考虑了。这是一种剪枝优化。

//maxHeap
struct cmp {
    bool operator()(pair<int, int> &a, pair<int, int> &b) {
        return a.first + a.second < b.first + b.second;
    }
};

class Solution {
public:
    /**
     * @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
     */
    int kthSmallestSum(vector<int> &a, vector<int> &b, int k) {
        int aSize = a.size();
        int bSize = b.size();
        priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> maxHeap;
        int count = 0;
        for (int i = 0; i < aSize; i++) {
            for (int j = 0; j < bSize; j++) {
                if (!maxHeap.empty() && count >= k && a[i] + b[j] > maxHeap.top().first + maxHeap.top().second) break;  //剪枝
                maxHeap.push({a[i], b[j]});
                count++;
                if (count > k) maxHeap.pop();
            }
        }
        return maxHeap.top().first + maxHeap.top().second;
    }
};

解法3： binary search

TBD。