面试题 17.14

Design an algorithm to find the smallest K numbers in an array.

Example:

Input: arr = [1,3,5,7,2,4,6,8], k = 4 Output: [1,2,3,4] Note:

0 <= len(arr) <= 100000 0 <= k <= min(100000, len(arr))

Solutions

1. qucik select (deduce and conquer) O(n + klog(k))

2. use stl's algorithm

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        nth_element(arr.begin(), arr.begin() + k, arr.end());
        return {arr.begin(), arr.begin() + k};
    }
};

• use home made version

#include <experimental/random>

class Solution {
public:
    template <typename It, typename Cmp = less<typename It::value_type>>
    void nth_element(It lo, It mid, It hi, Cmp && cmp = Cmp()) {
        typename iterator_traits<It>::difference_type zero{0};
        while (lo < hi) {
            iter_swap(lo, lo + std::experimental::randint(zero, hi - lo - 1));
            auto pivot = *lo;
            It st = lo, ed = hi - 1;
            while (st < ed) {
                while (st < ed && !cmp(*ed, pivot)) --ed;
                *st = *ed;
                while (st < ed && cmp(*st, pivot)) ++st;
                *ed = *st;
            }
            *st = pivot;
            // mid >= st, in either case, st can be passed
            if (!(mid < st)) lo = st + 1;
            // st >= mid, in either case, st can be passed
            if (!(st < mid)) hi = st;
        }
    }
    vector<int> smallestK(vector<int>& arr, int k) {
        nth_element(arr.begin(), arr.begin() + k, arr.end());
        return {arr.begin(), arr.begin() + k};
    }
};

1. heap O(nlog(k))

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        if (!k) return {};
        priority_queue<int> maxq;

        for (auto n : arr) {
            maxq.push(n);
            if (maxq.size() > k)
                maxq.pop();
        }

        vector<int> res(k);
        while (k--) {
            res[k] = maxq.top();
            maxq.pop();
        }

        return res;
    }
};