分治法解决选择问题——找出第i小的元素

给出一组数据，返回第i小的元素，最常规的思路是遍历找，或者排好序之后返回等，这里介绍用分治法解决，用到之前快排里的一个函数partition，为了把数组分成两部分，A[p,q-1]的元素都小于主元q， A[q+1,r]的元素都大于q， 所以只要比较i这个位置要落在哪一边就好了
伪代码如下：

int Select（int *A, int p, int r, int i）//对数组A[p,r]返回第i小的元素
{
    if(p==r)//只有一个元素
    return A[p];
    int q = Partition(A, p ,r);
    int k = q-p+1;//第一段的元素数
    if(i==k)//正好是划分的主元
    return A[q];
    else if(i<k)//落在前半部分
    return Select(A,p,q-1,i);
    else
    return Select(A,q+1,r,i-k);//落在后半部分，找第i-k个，因为前面已经k个了
}

算法期望是线性的O(n)，有重复元素也可用
完整代码如下:

#include<iostream>
#include<cstdio>
#include<algorithm>
using namespace std;
int Partition(int *A, int p ,int r)
{
    int x = A[r];
    int i = p-1;
    for(int j  = p;j<=r-1;j++)
    {
        if(A[j] <= x)
        {
            i++;
            swap(A[i],A[j]);
        }
    }
    swap(A[i+1],A[r]);
    return i+1;
}
int Select(int *A, int p, int r, int i)
{
    if(p==r)
        return A[p];
    int q = Partition(A,p,r);
    int k = q-p+1;
    if(i==k)
        return A[q];
    else if(i<k)
        return Select(A,p,q-1,i);
    else
        return Select(A,q+1,r,i-k);
}
int A[30000];
int main()
{
  int n,ii;
  while(cin>>n>>ii)
  {
      for(int i = 1;i<=n;i++)
        cin>>A[i];
    int ans = Select(A,1,n,ii);
    cout<<ans<<endl;
  }
}