Top K问题

Top K问题在数据分析中非常普遍的一个问题（在面试中也经常被问到），比如：

从20亿个数字的文本中，找出最大的前100个。

 

解决Top K问题有两种思路，

最直观：小顶堆（大顶堆 -> 最小100个数）；

较高效：Quick Select算法。

 

1. 堆

 

小顶堆（min-heap）有个重要的性质——每个结点的值均不大于其左右孩子结点的值，则堆顶元素即为整个堆的最小值。

小顶堆解决Top K问题的思路：小顶堆维护当前扫描到的最大100个数，其后每一次的扫描到的元素，若大于堆顶，则入堆，然后删除堆顶；依此往复，直至扫描完所有元素。

 

typedef multiset<int, std::less<int> > IntHeap; // 从小到大，小顶推
void FindTopKNumbers
(
      const vector<int>& data,               // a vector of data
      IntHeap& topNumbers,                 // top k, output
      unsigned int k                              
)
{
      topNumbers.clear();
      if(k == 0 || data.size() < k)
            return;
      vector<int>::const_iterator iter = data.begin();
      for(; iter != data.end(); ++ iter)
      {
            // if less than k numbers was inserted into topNumbers
            if((topNumbers.size()) < k)
                  topNumbers.insert(*iter);
            // topNumbers contains k numbers and it's full now
            else
            {
                  // first number in topNumbers is the least one
                  IntHeap::iterator iterFirst = topNumbers.begin();
                  // if is greater than the previous lease number
                  if(*iter > *(topNumbers.begin()))
                  {
                        // replace the previous least number
                        topNumbers.erase(iterFirst);
                        topNumbers.insert(*iter);
                  }
            }
      }
}

 

/*
  可以考虑维护一个堆，这个堆的中元素的个数为k，然后遍历元素，这时有两种情况：
    1、堆中元素的个数小于k，这时直接插入即可
    2、元素的个数为k，这时如果遍历到的元素小于堆中最大的那个元素就将堆中的这个元素删除，将遍历到的元素插入即可
*/
///
// find k least numbers in a vector
///
typedef multiset<int, std::greater<int> > IntHeap; // 从大到小，大顶推
//typedef multiset<int, std::less<int> > IntHeap; // 从小到大，小顶推
void FindKLeastNumbers
(
      const vector<int>& data,               // a vector of data
      IntHeap& leastNumbers,                 // k least numbers, output
      unsigned int k                              
)
{
      leastNumbers.clear();
      if(k == 0 || data.size() < k)
            return;
      vector<int>::const_iterator iter = data.begin();
      for(; iter != data.end(); ++ iter)
      {
            // if less than k numbers was inserted into leastNumbers
            if((leastNumbers.size()) < k)
                  leastNumbers.insert(*iter);
            // leastNumbers contains k numbers and it's full now
            else
            {
                  // first number in leastNumbers is the greatest one
                  IntHeap::iterator iterFirst = leastNumbers.begin();
                  // if is less than the previous greatest number
                  if(*iter < *(leastNumbers.begin()))
                  {
                        // replace the previous greatest number
                        leastNumbers.erase(iterFirst);
                        leastNumbers.insert(*iter);
                  }
            }
      }
}

2. Quick Select 

 

Quick Select 脱胎于快排（Quick Sort），两个算法的作者都是Hoare，并且思想也非常接近：选取一个基准元素pivot，将数组切分（partition）为两个子数组，比pivot大的扔左子数组，比pivot小的扔右子数组，然后递推地切分子数组。

Quick Select不同于Quick Sort的是其没有对每个子数组做切分，而是对目标子数组做切分。其次，Quick Select与Quick Sort一样，是一个不稳定的算法；pivot选取直接影响了算法的好坏，worst case下的时间复杂度达到了O(n2)。

 

int partition(int *v, const int low, const int high)
{
    if (v == NULL || low > high)
    {
        return -1;
    }
    int i = low;
    int j = high;
    const int pivot = v[i];
    while (i < j)
    {
        while (i < j && pivot > v[j])
        {
            --j;
        }
        if (i < j)
        {
            v[i++] = v[j];
        }
        while (i < j && pivot < v[i])
        {
            ++i;
        }
        if (i < j)
        {
            v[j--] = v[i];
        }
    }
    v[i] = pivot;
    return i;
}

int QuickSort(int *v, const int low, const int high)
{
    if (v == NULL || low < 0 || low > high)
    {
        return -1;
    }
    if (low < high)
    {
        int pivot = partition(v, low, high);
        QuickSort(v, low, pivot -1);
        QuickSort(v, pivot + 1, high);
    }
}

 

 

Quick Select的目标是找出第k大元素，所以

 

若切分后的左子数组的长度 > k，则第k大元素必出现在左子数组中；

若切分后的左子数组的长度 = k-1，则第k大元素为pivot；

若上述两个条件均不满足，则第k大元素必出现在右子数组中。

Quick Select的实现如下：

int QuickSelect(int *v, const int low, const int high, const int k)
{
    if (low > high || k < 1)
    {
        return -1;
    }
    if (low == high)
    {
        print(v, low);
        return v[low];
    }
    int pivot = partition(v, low, high);
    if (k == pivot - low + 1)
    {
        print(v, pivot);
        return v[pivot];
    }
    else if (k < pivot - low + 1)
    {
        QuickSelect(v, low, pivot - 1, k);
    }
    else
    {
        QuickSelect(v, pivot + 1, high, k - (pivot - low + 1));
    }
}

上面给出的代码都是求解第k大元素；若想要得到Top K元素，仅需要将代码做稍微的修改：比如，扫描完成后的小顶堆对应于Top K，Quick Select算法用中间变量保存Top K元素。