Friday, January 2, 2015

Find k-th smallest number in two sorted array

The algorithm is explained here.

I will explain it briefly based on my understanding. Choose any index in array A, say i. Choose a corresponding index in array B, say j,  that maintain the equivalence i + j = k - 1. Then the k-th smallest element can be acquired if one of the following conditions is fulfilled:

1. if k == 1, return the smallest between the first elements in A and B;
2. if (B[j - 1] <= A[i] < B[j +1]), then return A[i].
3. if (A[i - 1] <= B[j] < A[i] ), then return B[j].

Condition 2 and 3 are similar, take 2 as the example.






If A[i] is between B[j -1] and B[j], and there are i elements before A[i] and j elements before B[j], and because k = i + j + 1, then A[i] is the k-th smallest element we are looking for. 

Well, life would be perfect if one of the above conditions can always be fulfilled, what if not? 

So, if A[i] is not greater than or equal to B[j -1], but A[i] is smaller than B[j], what does that imply? It means A[i] must also be smaller than B[j - 1], and since k = i + j + 1, it must not be in A[i] and all elements less than A[i]. Then we can discard that part of the array. (Switch j - 1 and i in the figure to make it easier to understand). Likewise, we may conclude that all elements larger than B[j] will also not be in the range.  Similarly, we can discard corresponding elements if B[j] is smaller than A[i]. Good, we just halve the total array. 

And the solution will get by... Recursion! 

I follow the author of my reference and choose i and j  based on the weights of the lengths of two arrays. Alternatively, we can choose k/2 or anything else as long as both i and j are smaller than the lengths of A and B. 







public class FindKthSmallest {
 public int findKth(int[] A, int[] B, int k) {
  if (A == null || A.length == 0 || B == null || B.length == 0 
    || k < 0 || k > A.length + B.length)
   throw new IllegalArgumentException("Invalid input!");
  return findKthSmallest(A, 0, A.length, B, 0, B.length, k);
 }
 public int findKthSmallest(int[] A, int A_start, int la, int[] B, int B_start, int lb, int k) {
  if (A_start >= A.length)
   return B[B_start + k - 1];
  if (B_start >= B.length)
   return A[A_start + k - 1];
  if (k == 1)
   return Math.min(A[A_start], B[B_start]);
  
  //choose index based on the weight of length of A and B
  int A_key = (int) ((double) la / (la + lb) * (k - 1));
  int B_key = (k - 1) - A_key;
  //note the actual index will be A_start + A_key
  int index_A = A_start + A_key;
  int index_B = B_start + B_key;
  //index_A == 0: la = 1; lb = 10; k = 3; index_A = (1/11 *3) = 0
  int A_i_1 = (index_A == 0) ? Integer.MIN_VALUE : A[index_A - 1];
  int B_j_1 = (index_B == 0) ? Integer.MIN_VALUE : B[index_B - 1];
  //index_A == la: la = 10; lb = 1; k = 11; index_A = 10/11 * 11 = 10
  int A_i = (index_A == la) ? Integer.MAX_VALUE : A[index_A];
  int B_j = (index_B == lb) ? Integer.MAX_VALUE : B[index_B];
  
  if (B_j_1 <= A_i && A_i < B_j) 
   return A_i;
  if (A_i_1 <= B_j && B_j < A_i)
   return B_j;
  //if A_i < B_j and apparently from above A_i is NO larger than B_j_1
  if (A_i < B_j) 
   //discard A_i and below
   //length = la - (index_A - A_start + 1) = la - A_key -1
   return findKthSmallest(A, index_A + 1, la - A_key - 1, B, B_start, B_key + 1, k - A_key - 1);
  else
   return findKthSmallest(A, A_start, A_key + 1, B, index_B + 1, lb - B_key - 1, k - B_key - 1);
  
 }
}

No comments:

Post a Comment