程式員面試題精選（5）：尋找最小的k個元素

題目：輸入n個整數，輸出其中最小的k個。

例如輸入1，2，3，4，5，6，7和8這8個數字，則最小的4個數字為1，2，3和4。

分析：這道題最簡單的思路莫過於把輸入的n個整數排序，這樣排在最前面的k個數就是最小的k個數。只是這種思路的時間複雜度為O(nlogn)。我們試著尋找更快的解決思路。

我們可以開闢一個長度為k的數組。每次從輸入的n個整數中讀入一個數。如果數組中已經插入的元素少於k個，則將讀入的整數直接放到數組中。否則長度為k的數組已經滿了，不能再往數組裡插入元素，只能替換了。如果讀入的這個整數比數組中已有k個整數的最大值要小，則用讀入的這個整數替換這個最大值；如果讀入的整數比數組中已有k個整數的最大值還要大，則讀入的這個整數不可能是最小的k個整數之一，拋棄這個整數。這種思路相當於只要排序k個整數，因此時間複雜可以降到O(n+nlogk)。通常情況下k要遠小於n，所以這種辦法要優於前面的思路。

這是我能夠想出來的最快的解決方案。不過從給面試官留下更好印象的角度出發，我們可以進一步把代碼寫得更漂亮一些。從上面的分析，當長度為k的數組已經滿了之後，如果需要替換，每次替換的都是數組中的最大值。在常用的資料結構中，能夠在O(1)時間裡得到最大值的資料結構為最大堆。因此我們可以用堆（heap）來代替數組。

另外，自己重頭開始寫一個最大堆需要一定量的代碼。我們現在不需要重新去發明車輪，因為前人早就發明出來了。同樣，STL中的set和multiset為我們做了很好的堆的實現，我們可以拿過來用。既偷了懶，又給面試官留下熟悉STL的好印象，何樂而不為之？

參考代碼：

#include <set>
#include <vector>
#include <iostream>

using namespace std;

typedef multiset<int, greater<int> >  IntHeap;

///////////////////////////////////////////////////////////////////////
// find k least numbers in a vector
///////////////////////////////////////////////////////////////////////
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);
                  }
            }
      }
}