前 言

之前刷Leetcode題：最接近原點的 K 個點，本題直接用sort()排序會超時，看到大佬使用了一個叫nth_element()的函數(shù)，因此本文介紹STL中nth_element()的使用。

nth_element

首先看下函數(shù)原型：

template <class RandomAccessIterator>
void nth_element (RandomAccessIterator first, RandomAccessIterator nth,
                    RandomAccessIterator last);

函數(shù)說明：
重新排列range[first,last)中的元素，使第nth個位置的元素是按排序順序在該位置的元素。其他元素沒有任何特定的順序，只是第nth個元素之前的元素都不大于該元素，而第nth個元素后面的元素均小于該元素。
這個很好理解，舉個例子

//原始數(shù)列
[2, 8, 3, 7, 1, 9, 4, 5, 6, 0]

//nth_element nth = 5
[x, x, x, x, x, 5, x, x, x, x]

nth_element()將會使得第nth個位置的元素是按排序順序在該位置的元素，即為5，其他元素沒有任何特定的順序，但保證5之前的元素均比5小，5之后的元素均比5大。(類似于快速排序的原理)

當然，我們可以自定義判斷函數(shù)，作為nth_element()的第四個參數(shù)

template <class RandomAccessIterator, class Compare>
void nth_element (RandomAccessIterator first, RandomAccessIterator nth,
                    RandomAccessIterator last, Compare comp);

因此上述Leetcode題的解法：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        // 法一：排序(超時)

        // 法二：利用multimap
        // multimap<int, vector<int>> data;
        // vector<vector<int>> res;
        // int k = 0;

        // for (int i = 0; i < points.size(); ++i) {
        //     data.insert(pair<int, vector<int>>(pow(points[i][0], 2) + pow(points[i][1], 2), points[i]));
        // }

        // multimap<int, vector<int>>::iterator it = data.begin();
        // while(k < K) {
        //     res.push_back(it->second);
        //     it++;
        //     k++;
        // }

        // return res;

        //法三：優(yōu)先隊列 priority_queue
        // priority_queue<pair<int, vector<int>>> q;
        // int i;

        // for (i = 0; i < K; ++i) {
        //     q.push(pair<int, vector<int>>(pow(points[i][0], 2) + pow(points[i][1], 2), points[i]));
        // }

        // for (; i < points.size(); ++i) {
        //     int dis = pow(points[i][0], 2) + pow(points[i][1], 2);
        //     if (dis < q.top().first) {
        //         q.pop();
        //         q.push(pair<int, vector<int>>(dis, points[i]));
        //     }
        // }

        // vector<vector<int>> res;
        // while (!q.empty()) {
        //     res.push_back(q.top().second);
        //     q.pop();
        // }
        // return res;

        //法四：nth_element
        nth_element(points.begin(), points.begin() + K, points.end(), [](vector<int>& ptA, vector<int>& ptB){
            return ptA[0] * ptA[0] + ptA[1] * ptA[1] < ptB[0] * ptB[0] + ptB[1] * ptB[1];
        });
        return {points.begin(), points.begin()+K};
    }
};

也可以學習下其他的解法。

進一步討論

本人在VS 2019(MSVC 2017 編譯器)下實測，發(fā)現(xiàn)了一個問題
測試代碼：

int main()
{
    vector<int> numbers;
    for (int i = 1; i < 15; i++) numbers.push_back(i);
    random_shuffle(numbers.begin(), numbers.end()); //重新排列數(shù)組

    for (auto i : numbers) {
        cout << i << " ";
    }
    cout << endl;

    nth_element(numbers.begin(), numbers.begin() + 5, numbers.end());
    for (auto i : numbers) {
        cout << i << " ";
    }
}

運行結(jié)果：

結(jié)果

奇怪的事情發(fā)生了，為什么nth_element()把整個數(shù)組都排列了？
查看了下源碼，發(fā)現(xiàn)了其中的原因：

// FUNCTION TEMPLATE nth_element
template <class _RanIt, class _Pr>
_CONSTEXPR20 void nth_element(_RanIt _First, _RanIt _Nth, _RanIt _Last, _Pr _Pred) {
    // order Nth element, using _Pred
    _Adl_verify_range(_First, _Nth);
    _Adl_verify_range(_Nth, _Last);
    auto _UFirst     = _Get_unwrapped(_First);
    const auto _UNth = _Get_unwrapped(_Nth);
    auto _ULast      = _Get_unwrapped(_Last);
    if (_UNth == _ULast) {
        return; // nothing to do
    }

    while (_ISORT_MAX < _ULast - _UFirst) { // divide and conquer, ordering partition containing Nth
        auto _UMid = _Partition_by_median_guess_unchecked(_UFirst, _ULast, _Pass_fn(_Pred));

        if (_UMid.second <= _UNth) {
            _UFirst = _UMid.second;
        } else if (_UMid.first <= _UNth) {
            return; // Nth inside fat pivot, done
        } else {
            _ULast = _UMid.first;
        }
    }

    _Insertion_sort_unchecked(_UFirst, _ULast, _Pass_fn(_Pred)); // sort any remainder
}

源碼中的_ISORT_MAX為32，所以，在VS環(huán)境下，當數(shù)組長度小于32時，nth_element()會將數(shù)組全排序。數(shù)組超過32長度時，才會排序第nth位數(shù)字。
使用超過32長度的數(shù)組，再試一次，代碼都不展示了

結(jié)果

得到了預期的結(jié)果。

更多我的Leetcode題解，詳見Leetcode題解