c++ - 如何有效地找到数字流中元素的排名？

Question

最近我试图找到具有以下条件的数字流的中位数:

1. 3遍算法
2. O(nlog(n)) 时间
3. O(sqrt(n)) 空间

输入重复 3 次，其中 n(整数个数)，后跟 n 个整数 a_i，这样:

1. n 是奇数
2. 1≤n≤10^7
3. |a_i| ≤ 2^{30}

输入数据的格式如下所示:

5
1 3 4 2 5
5
1 3 4 2 5
5
1 3 4 2 5

到目前为止我的代码如下所示:

#ifdef STREAMING_JUDGE
#include "io.h"
#define next_token io.next_token
#else
#include<string>
#include<iostream>
using namespace std; 
string next_token()
{
    string s;
    cin >> s;
    return s;
}
#endif

#include<cstdio>
#include<cstdlib>
#include<vector>
#include<algorithm>
#include<iostream>
#include<math.h>

using namespace std;

int main()
{
    srand(time(NULL));
    //1st pass: randomly choose sqrt(n) numbers from the given stream of numbers
    int n = atoi(next_token().c_str());
    int p = (int)ceil(sqrt(n));
    vector<int> a;
    for(int i=0; i<n; i++)
    {
        int s=atoi(next_token().c_str());
        if( rand()%p == 0 && (int)a.size() < p )
        {
            a.push_back(s);
        }
    }
    sort(a.begin(), a.end());
    //2nd pass: find the k such that the median lies in a[k] and a[k+1], and find the rank of the median between a[k] and a[k+1]
    next_token();
    vector<int> rank(a.size(),0);
    for( int j = 0; j < (int)a.size(); j++ )
    {
        rank.push_back(0);
    }
    for( int i = 0; i < n; i++ )
    {
        int s=atoi(next_token().c_str());
        for( int j = 0; j < (int)rank.size(); j++ )
        {
            if( s<=a[j] )
            {
                rank[j]++;
            }
        }
    }
    int median = 0;
    int middle = (n+1)/2;
    int k;
    if( (int)a.size() == 1 && rank.front() == middle )
    {
        median=a.front();
        cout << median << endl;
        return 0;
    }
    for( int j = 0; j < (int)rank.size(); j++ )
    {
        if( rank[j] == middle )
        {
            cout << rank[j] << endl;
            return 0;
        }
        else if( rank[j] < middle && rank[j+1] > middle )
        {
            k = j;
            break;
        }
    }
    //3rd pass: sort the numbers in (a[k], a[k+1]) to find the median
    next_token();
    vector<int> FinalRun;
    if( rank.empty() )
    {
        for(int i=0; i<n; i++)
        {
            a.push_back(atoi(next_token().c_str()));
        }
        sort(a.begin(), a.end());
        cout << a[n>>1] << endl;
        return 0;
    }
    else if( rank.front() > middle )
    {
        for( int i = 0; i < n; i++ )
        {
            int s = atoi(next_token().c_str());
            if( s < a.front() )  FinalRun.push_back(s);
        }
        sort( FinalRun.begin(), FinalRun.end() );
        cout << FinalRun[middle-1] << endl;
        return 0;
    }
    else if ( rank.back() < middle )
    {
        for( int i = 0; i < n; i++ )
        {
            int s = atoi(next_token().c_str());
            if( s > a.back() )  FinalRun.push_back(s);
        }
        sort( FinalRun.begin(), FinalRun.end() );
        cout << FinalRun[middle-rank.back()-1] << endl;
        return 0;
    }
    else
    {
        for( int i = 0; i < n; i++ )
        {
            int s = atoi(next_token().c_str());
            if( s > a[k] && s < a[k+1] )  FinalRun.push_back(s);
        }
        sort( FinalRun.begin(), FinalRun.end() );
        cout << FinalRun[middle-rank[k]-1] << endl;
        return 0;
    }
}

但我仍然无法达到 O(nlogn) 时间复杂度。 我猜想瓶颈在于排名部分(即通过查找输入流中采样的 a[i] 的排名来查找 (a[k], a[k+1]) 中中位数的排名数字。)在第二遍。这部分在我的代码中具有 O(nsqrt(n)) 。

但是我不知道如何提高排名效率...... 有什么提高效率的建议吗？提前致谢!

对“rank”的进一步解释:采样数的rank计算流中小于或等于采样数的数字的数量。例如:在上面给出的输入中，如果对数字a[0]=2、a[1]=4和a[2]=5进行采样，则rank[0]=2，因为有两个数字( 1 和 2) 在流中小于或等于 a[0]。

感谢您的所有帮助。特别是 @alexeykuzmin0 的建议确实可以将第二遍加速到 O(n*logn) 时间。但还有一个问题:在第 1 遍中，我以 1/sqrt(n) 的概率对数字进行采样。当没有采样到的数字时(最坏的情况)， vector a为空，导致后面的遍无法执行(即发生段错误(核心转储))。 @Aconcagua，“选择所有剩余元素，如果没有超过所需的元素”是什么意思？谢谢。

Answer 1

你是对的，你的第二部分在 O(n√n) 时间内完成:

for( int i = 0; i < n; i++ )                    // <= n iterations
  ...
    for( int j = 0; j < (int)rank.size(); j++ ) // <= √n iterations

要解决这个问题，我们需要摆脱内部循环。例如，我们可以首先计算落入每个区间的数组元素的数量，而不是直接计算初始数组中小于阈值的元素数量:

// Same as in your code
for (int i = 0; i < n; ++i) {
    int s = atoi(next_token().c_str());
    // Find index of interval in O(log n) time
    int idx = std::upper_bound(a.begin(), a.end(), s) - a.begin();
    // Increase the rank of only that interval
    ++rank[idx];
}

然后计算阈值元素的排名:

std::partial_sum(rank.begin(), rank.end(), rank.begin());

最终的复杂度为O(n log n) + O(n) = O(n log n)。

---

这里我使用了两种STL算法:

1. std::upper_bound它使用二分查找方法在对数时间内查找排序数组中严格大于给定数字的第一个元素。
2. std::partial_sum它计算线性时间内给定数组的部分和。