Maximum Sum Combinations

Problem

Given two equally sized 1-D arrays A, B containing N integers each.

A sum combination is made by adding one element from array A and another element of array B.

Return the maximum C valid sum combinations from all the possible sum combinations.

Solution Approach

Expected Time complexity: $O(n*log(n))$

Click - to see solution code

C++

#define f f
#define s s

vector<int> Solution::solve(vector<int> &A, vector<int> &B, int C) {
    int n = A.size();
    vector<int> ans;
    sort(A.begin(), A.end(), greater<int>());
    sort(B.begin(), B.end(), greater<int>());

    priority_queue<pair<int, pair<int, int>>> pq;  // A[i]+B[j], <i, j>
    set<pair<int, int>> vis;                       // i, j

    pq.push({A[0] + B[0], {0, 0}});
    vis.insert({0, 0});

    for (int i = 0; i < C; ++i) {
        auto p = pq.top();
        pq.pop();
        ans.push_back(p.f);

        if (p.s.f + 1 < n && vis.find({p.s.f + 1, p.s.s}) == vis.end()) {
            vis.insert({p.s.f + 1, p.s.s});
            int x = A[p.s.f + 1] + B[p.s.s];
            pq.push({x, {p.s.f + 1, p.s.s}});
        }
        if (p.s.s + 1 < n && vis.find({p.s.f, p.s.s + 1}) == vis.end()) {
            vis.insert({p.s.f, p.s.s + 1});
            int x = A[p.s.f] + B[p.s.s + 1];
            pq.push({x, {p.s.f, p.s.s + 1}});
        }
    }

    return ans;
}