Finally I did it !!

A friend of mine referred me to a problem from my own blog (i.e. here :D) and I must say I enjoyed every aspect of the problem. From designing its awesum O(nlogn) solution to coding it with handling all the minor details. It’s about to be 6 A.M. now and I can’t sleep without posting the code 😀

Problem Statement: You are given 2 arrays A, B sorted in decreasing order of size m and n respectively and a number k, such that 1 <= k <= m * n. You need to fine the kth largest sum(a+b) possible where a in A, and b in B.

The only caveat is that I did it for kth smallest, I find ascending order to be more natural 🙂

#include<iostream>
#include<vector>

using namespace std;

#define MID(min, max) ((min) + (((max) - (min)) / 2))

int GetSmallerPairsCount(const vector<int>& A,
                         const vector<int>& B,
                         const int sum) {
  int num_smaller_pairs = 0;
  int temp = 0;
  // O(n) loop
  for (int A_index = A.size() - 1, B_index = 0; A_index >= 0; --A_index) {
    for (; B_index < B.size(); ++B_index) {
      if (A[A_index] + B[B_index] <= sum) {
        ++temp;
      } else {
        break;
      }
    }
    num_smaller_pairs += temp;
  }
  return num_smaller_pairs;
}

bool CheckIfSumExists(const vector<int>& A, const vector<int>& B, int sum) {
  // O(n) loop
  for (int a = A.size() - 1, b = 0; a >= 0 && b < B.size();) {
    if (A[a] + B[b] == sum) return true;
    if (A[a] + B[b] < sum) {
      ++b;
    } else {
      --a;
    }
  }
  return false;
}

int FindNextLargestSum(const vector<int>& A,
                       const vector<int>& B,
                       const int sum) {
  int new_sum, n_l = A.back() + B.back();
  // O(n) loop
  for (int A_index = A.size() - 1, B_index = 0;
      A_index >= 0 && B_index < B.size();) {
    new_sum = A[A_index] + B[B_index];
    if ((n_l > new_sum) && (new_sum > sum)) {
      n_l = new_sum;
    }
    if (new_sum <= sum) {
      ++B_index;
    } else {
      --A_index;
    }
  }
  return n_l;
}

int FindPreviousSmallestSum(const vector<int>& A,
                            const vector<int>& B,
                            const int sum) {
  int new_sum, p_s = A.front() + B.front();
  // O(n) loop
  for (int A_index = A.size() - 1, B_index = 0;
      A_index >= 0 && B_index < B.size();) {
    new_sum = A[A_index] + B[B_index];
    if ((p_s < new_sum) && (new_sum < sum)) {
      p_s = new_sum;
    }
    if (new_sum < sum) {
      ++B_index;
    } else {
      --A_index;
    }
  }
  return p_s;
}

// O(n*log(n)) order with O(1) extra space
int KthSmallestPairSum(const vector<int>& A,
                       const vector<int>& B,
                       const int k) {
  int max = A.back() + B.back();
  int min = A.front() + B.front();
  // O(log(n)) loop
  while (1) {
    int curr_mid = CheckIfSumExists(A, B, MID(min, max)) ? MID(min, max) : FindNextLargestSum(A, B, MID(min, max));
    int prev_mid = FindPreviousSmallestSum(A, B, curr_mid);
    int next_mid = FindNextLargestSum(A, B, curr_mid);

    int curr_mid_smaller = GetSmallerPairsCount(A, B, curr_mid);
    int prev_mid_smaller = GetSmallerPairsCount(A, B, prev_mid);
    int next_mid_smaller = GetSmallerPairsCount(A, B, next_mid);

    if (k == prev_mid_smaller) return prev_mid;
    if (prev_mid_smaller < k && next_mid_smaller >= k) {
      if (curr_mid_smaller < k) return next_mid;
      return curr_mid;
    }

    if (prev_mid_smaller > k) {
      max = FindPreviousSmallestSum(A, B, prev_mid);
    } else if (next_mid_smaller < k) {
      min = FindNextLargestSum(A, B, next_mid);
    } else {
      cout << "BUG\n";
      exit(0);
    }
  }
  return -2;
}

int main() {
  int size;
  cout << "Length of first array: " ;
  cin >> size;
  cout << "Enter the elements of first array\n";
  vector<int> A;
  for (int t = -1; size > 0; --size) {
    cin >> t;
    A.push_back(t);
  }
  cout << "Size of second array: ";
  cin >> size;
  vector<int> B;
  for (int t = -1; size > 0; --size) {
    cin >> t;
    B.push_back(t);
  }
  sort(A.begin(), A.end());
  sort(B.begin(), B.end());

  // Build the brute-force solution beforehand to compare
  // the results of our optimized method
  vector<int> brute;
  for (int i = 0; i < A.size(); ++i) {
    for (int j = 0; j < B.size(); ++j) {
      brute.push_back(A[i] + B[j]);
    }
  }
  sort(brute.begin(), brute.end());

  for (int kk = 1; kk <= A.size() * B.size(); ++kk) {
    int k = kk;
    if (k < 1 || k > A.size() * B.size()) {
      cout << "Invalid k. 1 <= k <= " << A.size() * B.size();
    } else {
      int result = KthSmallestPairSum(A, B, k);
      if (result != brute[k-1]) {
        cout << "Method broke. Wake up Avi !!";
        exit(0);
      }
      cout << k << "th smallest pair sum: " << result;
    }
    cout << endl;
  }
  cout << "All Iz Bhel !! " << endl;
  return 0;
}

Advertisements

About Avi Dullu
Jaako rakhe saiyan, maar sake na koi!

6 Responses to Finally I did it !!

  1. Hi, I found ur blog quite useful but I have one suggestion to u that instead of writing whole code if u can write, the logic that made it possible and link that logic to original theory and put the complete code in download area. It will make ur blog quite interesting.

    In DP problems, code are usually very simple but logic needs a good practice.

    Plz let me know if u like my comments…………….. 🙂

    • Avi Dullu says:

      Hey … thanx for the acknowledgement and comment 🙂

      BTW I generally only give (confusing) hints to the questions I post here but when I get a nice logic I like to code it myself. After all I am also a programmer 🙂

      • adf says:

        adf says:
        August 16, 2011 at 4:30 pm
        agreed that you would like to write code but it will be easy for anyone who stop by your blog to see the algo instead of your code … just a suggestion of course it is up to you

  2. Mahi says:

    can you please explain your algorithm? Its really tough to go through the whole program.

  3. Abhishek says:

    Really nice approach Avi. I think you can you can modify line 103 as max = prev_mid -1; and line 105 as min = next_mid + 1; It will save some extra O(n) calls.

  4. anton says:

    i request you to post logic/pseudocode instead of actual implementation, because understanding someone else’s code doesn’t help

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: