[Leetcode力扣 973] K Closest Points to Origin

一、题目描述

中等难度。给一堆二维点，要求输出距离原点(0, 0)最近的K个点

二、解法一

看到这道题，最容易想到的就是直接对整个vector从小到大排序，然后返回前K个元素即可，我们需要做的只是根据定义重写一个比较函数即可，时间复杂度为O(NlogN)。C++代码如下：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) 
    {
        sort(points.begin(), points.end(), cmp);
        return vector<vector<int>> (points.begin(), points.begin() + K);
    }
private:
    static bool cmp(const vector<int>& a, const vector<int>& b)
    {
        return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
    }
};

思考上面解法对解决这个问题的效率如何？很显然，题目只要求找到最小(也就是离原点最近)的K个元素，甚至不要求这些元素之间的顺序，而我们为了达到这个目的，把整个数组都排了序，显然做了很多冗余的工作，所以下面就要想办法优化解决办法。

三、解法二

其实STL中有一个叫做partial_sort()的函数，支持对部分数组进行排序。

这个函数接受三个迭代器作为参数，(在下文中用begin，middle和end代替)

partial_sort(begin_iterator, middle_iterator, end_iterator);

这个函数做的工作就是把左闭右开区间[begin, middle)中的元素以递增顺序排列，并保证这个区间内的元素是最小的(middle - begin)个，而不保证剩下的元素的顺序。把这个函数应用到本题中就是我们把前K个元素进行部分排序即可，C++代码如下：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) 
    {
        partial_sort(points.begin(), points.begin() + K, points.end(), cmp);
        return vector<vector<int>> (points.begin(), points.begin() + K);
    }
    
private:
    static bool cmp(const vector<int>& a, const vector<int>& b)
    {
        return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
    }
};

partial_sort()函数的原理是：将[begin, middle)区间内的元素建堆(如果元素想从小到大排的话建的是最大堆，所以以下均以最大堆为例)，然后对于[middle, end)范围内的元素逐一和最大堆的最大元素进行比较，如果小于最大堆的最大元素，则交换这两个元素，并重新组织整个最大堆。当最小的(middle - begin)个元素都被放到堆里去以后，再对整个堆进行排序。

所以这个解法的时间复杂度为O(NlogK + KlogK) = O(NlogK)。因为每次组织堆结构时间为O(logK)，一共进行了O(N)次，最后对整个堆进行堆排序的时间复杂度为O(KlogK)。

思考这个解法，时间效率优于解法一，因为我们只对K个较小元素进行了排序，但是细想想我们还是做了冗余工作，为了找到前K小的元素，我们对K个元素建了堆，还对前K个元素进行了堆排序，其实这些工作还是有点多余，因为题目根本不要求找到的元素需要满足什么结构或者顺序，那么还可以继续优化吗？可以！这就引出了下面的解法三。

四、解法三

其实STL中有一个叫做nth_element()的函数，这个函数也是接受三个迭代器，参照上文我们把三个迭代器也叫作begin，middle和end。这个函数的功能是将middle位置的元素变为排序后应该处在middle位置上的这个元素，并且所有middle左边的元素都比它小，右边的元素都比它大。这个算法的运行时间为O(N)，C++代码如下，注意参数迭代器的范围和上一个方法略有不同：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) 
    {
        nth_element(points.begin(), points.begin() + K - 1, points.end(), cmp);
        return vector<vector<int>> (points.begin(), points.begin() + K);
    }
    
private:
    static bool cmp(const vector<int>& a, const vector<int>& b)
    {
        return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
    }
};

其实这个函数的底层原理基本上是一种叫做快速选择的算法(和快速排序有点点类似)。事实上到这里我们已经找到了这个题目的最优解，提交以上代码，Leetcode显示运行时间为280ms，超过了91%的解答。

但是解法二和解法三都用了STL中现成的函数，如果在面试中面试官不让我们用STL中的函数怎么办？我们已经知道了nth_element()的原理是快速选择算法，那么我们手写一个快速选择算法就可以了，这就引出了解法四。

五、解法四：手写快速选择算法

快速选择算法和快速排序算法一样，关键在于partition()函数，该函数接受一个范围，在范围内选择一个pivot(可以随便选也可以随机选，出于性能角度考虑，我们把pivot选为范围内随机下标的元素)，该函数返回这个pivot代表的元素在排好序的数组中应该处于的位置，也就是说这个位置左边的元素都小于等于它，右边都大于等于它。

基于上述partition()函数，我们就可以得出快速选择算法的逻辑：当partition()返回的pivotIndex == K-1时，说明已经找到了最小的K个元素；当pivotIndex > K-1时，说明pivot跑到了"K的右边"，这时应该在pivotIndex左边寻找；当pivotIndex < K-1时，说明pivot跑到了“K的左边”，这时应该在pivotIndex的右边寻找。

使用上面快速选择算法的逻辑，解决本题的C++代码如下：

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) 
    {
        int left = 0, right = points.size() - 1;
        while(true)
        {
            int pivotIndex = partition(points, left, right);
            if(K - 1 == pivotIndex)
                break;
            else if(K - 1 > pivotIndex)
                left = pivotIndex + 1;
            else
                right = pivotIndex - 1;
        }
        return vector<vector<int>> (points.begin(), points.begin() + K);
    }
private:
    bool isless(const vector<int>& a, const vector<int>& b)
    {
        return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
    }
    
    bool isgreater(const vector<int>& a, const vector<int>& b)
    {
        return a[0]*a[0] + a[1]*a[1] > b[0]*b[0] + b[1]*b[1];
    }
    
    int partition(vector<vector<int>>& points, int left, int right)
    {
        if(left == right)
            return left;
        
        srand(time(0));
        int randomIndex = rand() % (right - left + 1) + left;
        swap(points[randomIndex], points[left]);
        
        vector<int> pivot = points[left];
        while(left < right)
        {
            while(left < right && !isless(points[right], pivot))
                --right;
            points[left] = points[right];
            while(left < right && !isgreater(points[left], pivot))
                ++left;
            points[right] = points[left];
        }
        points[left] = pivot;
        return left;
    }
};