求 n + 2 个数的平均数

题目描述

给定一个按非递减排序的正整数数组𝐴[1,2, … , 𝑛]，以及两个正整数 𝑥 和 𝑦，请设计算法求这 𝑛 + 2 数的中位数。这里的中位数定义为 𝑛 + 2 为奇数时，中位数为最中间大小的那个数， 𝑛 + 2 为偶数时，中位数为最中间大小的两个数的均值，并分析算法的时间代价。

问题分析

两个正整数，可以将较小的数记为 x，较大的数记为 y，方便后续的处理。

可以先分别查找原始数组中，不小于 x、不小于 y 的数的下标，分别记为 px 和 py。

若数组个数为奇数，即中位数是中间的那个数，将原数组中中位数的下标记为 idx：

• 若px <= idx && py >= idx：中位数不变，仍是a[idx]
• 若py < idx：中位数可能是y或a[idx-1]，取二者最大值。
• 若px > idx：中位数可能是x或a[idx+1]，取二者最小值。

若数组个数为偶数，即中位数是最中间的两个数的均值，将这两个数的下标分别记为 idx - 1 和 idx：

• 若px <= idx-1 && py >= idx：中位数为(a[idx-1] + a[idx]) / 2.0
• 若py <= idx-1：中位数为(max(y, a[idx-2])+ a[idx-1]) / 2.0
• 若px >= idx：中位数为(a[idx] + min(x, a[idx+1])) / 2.0

复杂度分析

时间开销主要在两次二分查找，为 $O(logn)$

程序代码

#include <iostream>
#include <algorithm>
#include <vector>
using namespace std;

// O(logn)
// 找到数组中大于等于num的第一个位置
int binarySearch(const vector<int>& a, int num)
{
    int l = 0, r = a.size();
    while(l < r) {
        int mid = (l + r) >> 1;
        if(a[mid] < num) {
            l = mid + 1;
        }
        else {
            r = mid;
        }
    }
    return l;
}

int main()
{
    int n;
    cin >> n;
    vector<int> a(n);
    for(int i = 0; i < n; i++) {
        cin >> a[i];
    }
    int x, y;
    cin >> x >> y;
    
    // 小的数作为x
    if(x > y)  swap(x, y);
    
    int px = binarySearch(a, x), py = binarySearch(a, y);

    // 中位数位置初始值
    int idx = n >> 1;
    double res;
    if( n & 1 ) {
        if(px <= idx && py >= idx) {
            res = a[idx];
        }
        else if( py < idx ) {
            res = max(y, a[idx-1]);
        }
        else {
            res = min(x, a[idx+1]);
        }
    }
    else {
        if(px <= idx-1 && py >= idx) {
            res = (a[idx-1] + a[idx]) / 2.0;
        }
        else if( py <= idx-1 ) {
            res = (max(y, a[idx-2])+ a[idx-1]) / 2.0;
        }
        else {
            res = (a[idx] + min(x, a[idx+1])) / 2.0;
        }
    }
    
    cout << res << endl;
    
    return 0;
}