xiaobaoqiu Blog

Think More, Code Less

两个已序数组的第k小元素

Problem:

两个已序数组(从小到大)的第k小元素.数组分别记为A(长度m)和B(长度n).

Solutions

方法1

将两个数组merge成一个有序数组(m+n大小),再直接取第k个元素.

时间复杂度:O(m+n).

空间复杂度:O(m+n).

方法2

两个指针i和j分别从小到大遍历两个数组,代码大致如下:

1
2
3
4
5

i=0, j=0
while((i+j)<k && (i<m || j<n)){
    if(A[i] < B[j]) i++;
    else j++;
}

时间复杂度:O(k).

空间复杂度:O(1).

方法3

利用方法2,我们发现,如果

1

B[j-1] < A[i] < B[j]

则,A[i]为第i+j+1小的元素,同理,如果

1

A[i-1] < B[j] < A[i]

则B[j]是第i+j+1小的元素.令

1
2
3

i+j+1 = k
即
i+j=k-1

则找到了我们需要的第k小元素.

下面是我们的方法:

  1. 二分搜索数组A,i的初值: i = (int)((double)m / (m+n) * (k-1)); 令j = k-1-i;

  2. 比较,如果 B[j-1] < A[i] < B[j] 则返回A[i]. 如果 A[i-1] < B[j] < A[i] 则返回B[j];

  3. 步骤2未结束,则比较A[i]和B[j]

    如果A[i]<B[j],则我们可以排除比A[i]小的元素和比B[j]大的元素,即

    在A[i+1, m]和B[0, j]中查找第k-i小的元素

    否在,可以排除比B[j]小的元素和比A[i]大的元素,即

    在A[0, i]和B[j+1, n]中查找第k-j小的元素

简单示例代码(不完整):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

int kthSmallest(int A[], int m, int B[], int n, int k) {
  
  int i = (int)((double)m / (m+n) * (k-1));
  int j = (k-1) - i;

  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;
 
  if (Ai < Bj)
    return kthSmallest(A+i+1, m-i-1, B, j, k-i);
  else
    return kthSmallest(A, i, B+j+1, n-j-1, k-j);
}

时间复杂度:O(logm + logn)

空间复杂度:O(1)