sundeepblue
5/15/2014 - 7:16 PM

Kth maximal element of sum of two arrays

Kth maximal element of sum of two arrays

// http://stackoverflow.com/questions/5212037/find-the-pair-across-2-arrays-with-kth-largest-sum

// result may not be correct!!!!!!!!!!!!!!!!!

// 6/6/2015. attempted to calculate time complexity seems to be:
// log1 + log2 + ... + logk = logk! = O(klogk)

// here is how to calculate logk!
// http://stackoverflow.com/questions/2095395/is-logn-%CE%98n-logn
// http://ballinger.disted.camosun.bc.ca/math126/lognfactorial.pdf

struct node {
    int i, j;
    int sum;
    node(int _i, int _j, int s) : i(_i), j(_j), sum(s) {}
};
struct comp {
    bool operator()(const node& n1, const node& n2) {
        return n1.sum < n2.sum;
    }  
};

pair<int,int> get_pair(vector<int> A, vector<int> B, int k) {
    sort(A.begin(), A.end());
    sort(B.begin(), B.end());
    pair<int, int> res = make_pair(-1, -1);
    if(k > A.size() + B.size()) return res;
    
    priority_queue<node, vector<node>, comp> q;
    q.push(node(0, 0, A[0] + B[0]));
    while(!q.empty()) {
        node nd = q.top();
        q.pop();
        k--;
        if(k == 0) return make_pair(nd.i, nd.j);
        int i = nd.i, j = nd.j;
        
        if(i< A.size() - 1) q.push(node(i+1, j, A[i+1] + B[j]));
        if(j < B.size() - 1) q.push(node(i, j+1, A[i] + B[j+1]));
    }
    return res;
}


int main()
{
   vector<int> a = {1,3,9};
   vector<int> b = {1,5};
   pair<int,int> res = get_pair(a, b, 5);
   cout << res.first << ", " << res.second;
   return 0;
}