在两个排序数组中找到第k小的数

在两个排序数组中找到第k小的数

题目描述

给定两个有序数组arr1和arr2，再给定一个整数K，返回所有数中第K小的数。

[要求]

如果arr1的长度为N，arr2的长度为M，时间复杂度请达到$O(\log (\min N,M))$，额外空间复杂度$O(1)$。

输入描述：

第一行三个整数N, M, K分别表示数组arr1, arr2的大小，以及需要询问的数
接下来一行N个整数，表示arr1内的元素
再接下来一行M个整数，表示arr2内的元素

输出描述：

输出一个整数表示答案

示例1

输入

5 3 1
1 2 3 4 5
3 4 5

输出

1

说明

1是所有数中第一小的数

示例2

输入

3 4 4
1 2 3
3 4 5 6

输出

3

说明

3是所有数中第4小的数，所以返回3

备注：

$1⩽N,M⩽10^5$
$0⩽ar{r}_{1_i},ar{r}_{2_i}⩽10^9$

---

题解：

在两个长度相等的排序数组中找到上中位数 进阶版题目，我们仍然可以按照这题的思想去做。

假设长度较短的数组为 shortArr，长度记为 lenS ；长度较长的数组为 longArr ，长度记为 lenL 。那么对于整数k，有以下三种情况：

• k <= lenS，这时，在 shortArr 和 longArr 中同时选择前 k 个数，两段数组中上中位数就是整体的第 k 小；
• k > lenL ，这时，分情况讨论：
  1. 首先 longArr 中的前 k-lenS-1 个元素是不符合的，因为即使它们都比 shortArr[lenS-1] 大，也不能到达第 k 个。若 longArr[k-lenS-1] >= shortArr[lenS-1] ，那么 longArr[k-lenS-1] 就是第 k 小，否则的话，它也不是；
  2. 其次shortArr 中的前 k-lenL-1 个元素是不符合的，因为即使它们都比 longArr[lenL-1] 大，也不能达到第 k 个。若 shortArr[k-lenL-1] >= longArr[lenL-1] ，那么 shortArr[k-lenL-1] 就是第 k 小，否则的话，它也不是；
  3. 在上述两个条件都不满足的情况下，shortArr[k-lenL…lenS-1] 和 longArr[k-lenS…lenL-1] 这两段数组的上中位数就是整体的第 k 小。
• lenS < k <= lenL，分情况讨论：
  1. 首先 longArr 中的前 k-lenS-1 个元素是不符合的，因为即使它们都比 shortArr[lenS-1] 大，也不能到达第 k 个。若 若 longArr[k-lenS-1] >= shortArr[lenS-1] ，那么它就是第 k 小，否则的话，它也不是。同时， longArr 数组中的第 k 个元素以后也不符合要求；
  2. 条件1不满足的话，shortArr[0…lenS-1] 和 longArr[k-lenS…k-1] 这两段数组的上中位数就是整体的第 k 小

代码：

#include <cstdio>
#include <algorithm>

using namespace std;

const int N = 100000;

int n, m, k;
int a[N];
int b[N];

int getUpMedian( int s1, int e1, int s2, int e2 ) {
    int m1, m2;
    int bit;
    while ( s1 < e1 ) {
        m1 = s1 + e1 >> 1;
        m2 = s2 + e2 >> 1;
        bit = (( e1 - s1 + 1 ) & 1) ^ 1;
        if ( a[m1] > b[m2] ) {
            e1 = m1;
            s2 = m2 + bit;
        } else if ( a[m1] < b[m2] ) {
            s1 = m1 + bit;
            e2 = m2;
        } else return a[m1];
    }
    return min( a[s1], b[s2] );
}

int findKth() {
    if ( k < 1 || k > n + m ) return -1;
    if ( k <= n ) return getUpMedian(0, k - 1, 0, k - 1);
    if ( k > m ) {
        if ( a[k - m - 1] >= b[m - 1] )
            return a[k - m - 1];
        if ( b[k - n - 1] > a[n - 1] )
            return b[k - n - 1];
        return getUpMedian(k - m, n - 1, k - n, m - 1);
    }
    if ( b[k - n - 1] >= a[n - 1] ) 
        return b[k - n - 1];
    return getUpMedian(0, n - 1, k - n, k - 1);
}

int main(void) {
    scanf("%d%d%d", &n, &m, &k);
    for ( int i = 0; i < n; ++i ) 
        scanf("%d", a + i);
    for ( int i = 0; i < m; ++i )
        scanf("%d", b + i);
    // 始终让 a 指向元素个数少的数组，方便后面的处理
    if ( n > m ) {
        swap( a, b );
        swap( n, m );
    }
    printf("%d\n", findKth());
    return 0;
}