Kth Element of Two Sorted Arrays

91.3k
0

Given two sorted arrays A and B of size M and N respectively. Find the kth element of the final sorted array. The final array is formed by merging both arrays and sorting them in non-decreasing order. k is one based index.

Example 1

Input: a = [2, 3, 6, 7, 9], b = [1, 4, 8, 10], k = 5

Output: 6

Explanation: If we merge the two sorted arrays, the combined sorted array will be [1, 2, 3, 4, 6, 7, 8, 9, 10]. The 5th element in this combined array is 6.

Example 2

Input: a = [100, 112, 256, 349, 770], b = [72, 86, 113, 119, 265, 445, 892], k = 7

Output: 256

Explanation: The combined sorted array is [72, 86, 100, 112, 113, 119, 256, 265, 349, 445, 770, 892]. The 7th element in this combined array is 256.

Brute Force Approach

The most direct thought is to put both arrays into one list, sort that list, and return the element at index k - 1.

This works because the problem itself talks about the final merged sorted order. Once that order is available, the answer is just one lookup.

The drawback is that the arrays are already sorted, but this approach does not use that advantage properly. It still sorts everything again.

Algorithm

  • Create an empty list to store all elements from both arrays. This is needed because the kth element belongs to the combined order, not just one array.

  • Add every element from the first array and then every element from the second array.

  • Sort the combined list so the final order becomes the same as the merged sorted array.

  • Return the value at index k - 1 because k is 1-based while array indexes are 0-based.

Dry Run

Kth Element of 2 Sorted Arrays Brute Dry Run

Kth Element of 2 Sorted Arrays Brute Dry Run

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
/*
Returns the kth element after combining
both sorted arrays into one sorted order.
*/
int kthElement(vector<int>& a, vector<int>& b, int k) {
vector<int> merged;
for (int value : a) {
merged.push_back(value);
}
for (int value : b) {
merged.push_back(value);
}
// Sorting gives the exact final order that the problem is asking about.
sort(merged.begin(), merged.end());
// k is 1-based, so the needed index is k - 1.
return merged[k - 1];
}
};
// Driver code starts
int main() {
vector<int> a = {2, 3, 6, 7, 9};
vector<int> b = {1, 4, 8, 10};
int k = 5;
Solution obj;
cout << obj.kthElement(a, b, k) << endl;
return 0;
}

Complexity Analysis

Time Complexity: O((N + M) x log(N + M)), because all elements are sorted together.

Space Complexity: O(N + M), because a separate merged list is created.

Better Approach

Since both arrays are already sorted, there is no need to sort again.

The merge step of merge sort gives a useful idea: compare the current elements of both arrays and pick the smaller one first.

But there is an even smaller observation. The answer is only the kth element, so the full merged array is not needed. It is enough to simulate the merge process only until k elements have been picked.

The last picked element at that moment is the answer.

Algorithm

  • Keep one pointer at the start of array A and one pointer at the start of array B. These pointers show the next unused element in each array.

  • Keep a counter to track how many elements have been picked from the imaginary merged array.

  • Compare A[i] and B[j]. Pick the smaller value because it must appear earlier in the sorted merged order.

  • After picking a value, increase the counter. If the counter becomes k, return that picked value because the kth position has just been reached.

  • If one array finishes first, continue picking from the other array because its remaining elements are already sorted.

Dry Run

Better

Better

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
/*
Returns the kth element by simulating
only the needed part of the merge process.
*/
int kthElement(vector<int>& a, vector<int>& b, int k) {
// i points to the next unused element in the first array.
int i = 0;
// j points to the next unused element in the second array.
int j = 0;
// count tells how many elements have been picked so far.
int count = 0;
// This stores the most recently picked value from the merged order.
int answer = -1;
while (i < (int)a.size() && j < (int)b.size()) {
// Pick from a when its current value is smaller or equal, because it comes next.
if (a[i] <= b[j]) {
answer = a[i];
i++;
} else {
// Pick from b because its current value is smaller than a[i].
answer = b[j];
j++;
}
count++;
// The last picked value is the answer as soon as k elements are picked.
if (count == k) {
return answer;
}
}
while (i < (int)a.size()) {
answer = a[i];
i++;
count++;
// If b is finished, the kth value may still be in the remaining part of a.
if (count == k) {
return answer;
}
}
while (j < (int)b.size()) {
answer = b[j];
j++;
count++;
// If a is finished, the kth value may still be in the remaining part of b.
if (count == k) {
return answer;
}
}
return -1;
}
};
// Driver code starts
int main() {
vector<int> a = {2, 3, 6, 7, 9};
vector<int> b = {1, 4, 8, 10};
int k = 5;
Solution obj;
cout << obj.kthElement(a, b, k) << endl;
return 0;
}

Complexity Analysis

Time Complexity: O(K), because only the first k elements of the merged order are simulated.

Space Complexity: O(1), because constant space is used.

Optimal Approach

The key observation is that the kth element is the largest element among the first k elements of the combined sorted array.

Imagine the merged array split into two parts:

First k elements | Remaining elements
       Left      |      Right

The answer is the largest value in the left part.

The first k elements can contain some elements from A and the remaining elements from B. Suppose cutA elements are taken from A. Then cutB = k - cutA elements must be taken from B.

For a correct partition, every element in the left part must be <= every element in the right part. Only the elements immediately around the two cuts need to be compared:

  • leftA → the last element taken from A

  • rightA → the first element not taken from A

  • leftB → the last element taken from B

  • rightB → the first element not taken from B

The partition is valid when:

leftA <= rightB
leftB <= rightA

When the partition is valid, the largest element among the first k elements is:

max(leftA, leftB)

Binary search is used to find the correct cutA. The smaller array is chosen for binary search so that the search range remains as small as possible.

Algorithm

  • Check whether k is valid. Return -1 if k is outside the range 1 to n + m.

  • Always perform binary search on the smaller array.

  • Set the possible range of cutA from max(0, k - m) to min(k, n).

  • Choose cutA and calculate cutB = k - cutA. Together, both cuts place exactly k elements in the left part.

  • Find the four boundary values around the cuts: the last selected and first unselected element from both arrays.

  • If leftA <= rightB and leftB <= rightA, the partition is correct. Return max(leftA, leftB).

  • If leftA > rightB, too many elements have been taken from A, so move the partition in A to the left.

  • Otherwise, too few elements have been taken from A, so move the partition in A to the right.

  • Use INT_MIN and INT_MAX when a partition reaches the beginning or end of an array, since no actual element exists on that side.

Dry Run

Optimal

Optimal

Solution

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
/*
Returns the kth element by finding a valid
partition between the two sorted arrays.
*/
int kthElement(vector<int>& a, vector<int>& b, int k) {
int n = (int)a.size();
int m = (int)b.size();
// Invalid k cannot point to any position in the combined sorted array.
if (k < 1 || k > n + m) {
return -1;
}
// Search the smaller array so the binary search stays short and safe.
if (n > m) {
return kthElement(b, a, k);
}
// cutA cannot take less than k - m elements, or cutB would become too large.
int low = max(0, k - m);
// cutA cannot take more than n elements, or more than k total left elements.
int high = min(k, n);
while (low <= high) {
// cutA decides how many elements are taken from the first array.
int cutA = low + (high - low) / 2;
// cutB fills the remaining left-side positions so total left count is k.
int cutB = k - cutA;
// If cutA is 0, no value exists on the left side of a.
int leftA = (cutA == 0) ? INT_MIN : a[cutA - 1];
// If cutA is n, no value exists on the right side of a.
int rightA = (cutA == n) ? INT_MAX : a[cutA];
// If cutB is 0, no value exists on the left side of b.
int leftB = (cutB == 0) ? INT_MIN : b[cutB - 1];
// If cutB is m, no value exists on the right side of b.
int rightB = (cutB == m) ? INT_MAX : b[cutB];
// Both left parts fit before both right parts, so the partition is correct.
if (leftA <= rightB && leftB <= rightA) {
return max(leftA, leftB);
} else if (leftA > rightB) {
// Too many elements were taken from a, so move cutA to the left.
high = cutA - 1;
} else {
// Too few elements were taken from a, so move cutA to the right.
low = cutA + 1;
}
}
return -1;
}
};
// Driver code starts
int main() {
vector<int> a = {2, 3, 6, 7, 9};
vector<int> b = {1, 4, 8, 10};
int k = 5;
Solution obj;
cout << obj.kthElement(a, b, k) << endl;
return 0;
}

Complexity Analysis

Time Complexity: O(log(min(N, M))), because binary search is performed solely on the smaller array where every iteration requires only constant-time O(1) operations to compute partitions and compare boundary elements.

Space Complexity: O(1), because constant space is used.

Interview follow-up Questions

If k is larger than the size of array A, we cannot pick 0 elements from array A, because even picking all elements from B wouldn't give us k elements. Similarly, if k is smaller than the size of A, we cannot pick more than k elements from A. Pointers prevent going out of bounds.

Two PointerSortingMathsHeapBinary SearchArrays

Read Similar Blogs

Arrays

Row wise traversal

Input: mat = [[1, 2, 3], [4, 5, 6]] Output: 1 2 3 4 5 6 Explanation: First row prints as 1 2 3 , then second row prints as 4 5 6 . Input: mat = [[10, 20], [30, 40], [50, 60]] Output: 10 20 30 40 50...

61.7k
Arrays

Main Diagonal Print

Given a matrix with N rows and M columns, print all elements present on the main diagonal. Input: mat = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] Output: 1 5 9 Explanation: Main diagonal cells are mat[0][0...

67.1k
Arrays

Search Element in Matrix

Given a matrix with N rows and M columns and a target value X , search for X in the matrix. Return the position of the first occurrence of X as {row, col} . Return {-1, -1} when the t...

98.5k
Comments0