Simple C++ Solution using Max Heap in O(nlogn) time


#1

My code follows the following approach :

  1. Sort both arrays array A and array B.
  2. Create a max heap i.e priority_queue in C++ to store the sum combinations along with the indices of elements from both arrays A and B which make up the sum. Heap is ordered by the sum.
  3. Initialize the heap with the maximum possible sum combination i.e (A[N – 1] + B[N – 1] where N is the size of array) and with the indices of elements from both arrays (N – 1, N – 1). The tuple inside max heap will be (A[N-1] + B[N – 1], N – 1, N – 1). Heap is ordered by first value i.e sum of both elements.
  4. Pop the heap to get the current largest sum and along with the indices of the element that make up the sum. Let the tuple be (sum, i, j).
  5. Next insert (A[i – 1] + B[j], i – 1, j) and (A[i] + B[j – 1], i, j – 1) into the max heap but make sure that the pair of indices i.e (i – 1, j) and (i, j – 1) are not. This can be checked by using a map of pair which would check that a particular pair (i,j) was already present in the heap or not by considering pair -> bool (key-pair) in it.
  6. Go back to Step 4 till our ans array has size C.
#define mp make_pair
#define pb push_back
vector<int> Solution::solve(vector<int> &A, vector<int> &B, int C)
{
  priority_queue<pair<int,pair<int,int> > >  heap; // max heap to store the sum along with the indexes from which the sum came from
  sort(A.rbegin(),A.rend()); // sorting array A in descending order
  sort(B.rbegin(),B.rend()); // sorting array B in descending order
  map<pair<int,int>,bool > m; // map to be used as a vis array whether index (i,j)'s sum is already present int the heap or not
  heap.push(mp(A[0]+A[0],mp(0,0))); // Initially pushing the sum of the biggest elements from both the arrays (Ofcourse it would be the biggest sum combination for our solution)
  m[mp(0,0)]=true; // this represent that sum from (0,0) is already present in the heap
  vector<int> ans;
  while(ans.size()<C) // now find the top elements and keep pushing the (i+1,j) & (i,j+1)th sum cobinations in the heap (remeber that the top of the heap would only contain the maximum sum)
  {
      pair<int,pair<int,int> > temp = heap.top();
      heap.pop();
      int sum = temp.first;
      ans.pb(sum);
      int ind1 = temp.second.first;
      int ind2 = temp.second.second;
      if(m[mp(ind1+1,ind2)]!=true){
          heap.push(mp(A[ind1+1]+A[ind2],mp(ind1+1,ind2)));
          m[mp(ind1+1,ind2)] = true;
      }
      if(m[mp(ind1,ind2+1)]!=true){
          heap.push(mp(A[ind1]+A[ind2+1],mp(ind1,ind2+1)));
          m[mp(ind1,ind2+1)] = true;
      }
      
  }
  return ans;
}