【ATT】Find Kth smallest element in Two sorted array

A: 二分

Maintaining the invariant
i + j = k – 1,

If B_j-1 < A_i < B_j, then A_i must be the k-th smallest,
or else if A_i-1 < B_j < A_i, then B_j must be the k-th smallest.

 

If one of the above conditions are satisfied, we are done. If not, we will use i and j as the pivot index to subdivide the arrays. But how? Which portion should we discard? How about A_i and B_j itself?

We make an observation that when A_i < B_j, then it must be true that A_i < B_j-1. On the other hand, if B_j < A_i, then B_j < A_i-1. Why?

Using the above relationship, it becomes clear that when A_i < B_j, A_i and its lower portion could never be the k-th smallest element. So do B_j and its upper portion. Therefore, we could conveniently discard A_i with its lower portion and B_j with its upper portion.

int findKthSmallest(int A[], int m, int B[], int n, int k) {
  assert(m >= 0); assert(n >= 0); assert(k > 0); assert(k <= m+n);
  
  int i = (int)((double)m / (m+n) * (k-1));
  int j = (k-1) - i;
 
  assert(i >= 0); assert(j >= 0); assert(i <= m); assert(j <= n);
  // invariant: i + j = k-1
  // Note: A[-1] = -INF and A[m] = +INF to maintain invariant
  int Ai_1 = ((i == 0) ? INT_MIN : A[i-1]);
  int Bj_1 = ((j == 0) ? INT_MIN : B[j-1]);
  int Ai   = ((i == m) ? INT_MAX : A[i]);
  int Bj   = ((j == n) ? INT_MAX : B[j]);
 
  if (Bj_1 < Ai && Ai < Bj)
    return Ai;
  else if (Ai_1 < Bj && Bj < Ai)
    return Bj;
 
  assert((Ai > Bj && Ai_1 > Bj) || 
         (Ai < Bj && Ai < Bj_1));
 
  // if none of the cases above, then it is either:
  if (Ai < Bj)
    // exclude Ai and below portion
    // exclude Bj and above portion
    return findKthSmallest(A+i+1, m-i-1, B, j, k-i-1);
  else /* Bj < Ai */
    // exclude Ai and above portion
    // exclude Bj and below portion
    return findKthSmallest(A, i, B+j+1, n-j-1, k-j-1);
}