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

## Recent Activity