PAT (Advanced Level) Practice-1029 Median

二分查找

#include<bits/stdc++.h>
using namespace std;

const int maxn = 2e5+5;
long long a[maxn], b[maxn];

long long getKth(int k, int lo1, int hi1, int lo2, int hi2) {
    if (lo1 == hi1) return b[lo2+k-1];
    if (lo2 == hi2) return a[lo1+k-1];
    if (k == 1) return min(a[lo1], b[lo2]);
    int mi1 = min(hi1, lo1 + (k >> 1)), mi2 = min(hi2, lo2 + (k >> 1));
    if (a[mi1-1] < b[mi2-1]) return getKth(k-mi1+lo1, mi1, hi1, lo2, hi2);
    else return getKth(k-mi2+lo2, lo1, hi1, mi2, hi2);
}
int main() {
    int na, nb;
    scanf("%d", &na); for (int i = 0; i < na; ++i) scanf("%lld", a + i);
    scanf("%d", &nb); for (int i = 0; i < nb; ++i) scanf("%lld", b + i);
    if ((na + nb) & 1) printf("%lld\n", getKth((na+nb)/2+1, 0, na, 0, nb));
    else printf("%lld\n", getKth((na+nb)/2, 0, na, 0, nb));
    return 0;
}