# Finally I did it !!

July 21, 2011 6 Comments

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; }

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…………….. 🙂

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:

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

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

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.

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