973

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        // Add distance from origin to end of each point.
        for(int i=0; i<points.size(); i++){
            int dist = points[i][0]*points[i][0] + points[i][1]*points[i][1];
            points[i].push_back(dist);
        }

        // Do quickselect.
        quickselect(0, points.size()-1, k-1, points);

        // Add first k points to result.
        vector<vector<int>> res;
        for(int i=0; i<k; i++)
            res.push_back({points[i][0], points[i][1]});

        return res;
    }

    void quickselect(int left, int right, int k, vector<vector<int>> &nums){
        if(left == right) return;
        int pivot = partition(left, right, nums);
        if(pivot == k)
            return;
        else if(pivot > k)
            quickselect(left, pivot-1, k, nums);
        else
            quickselect(pivot+1, right, k, nums);
    }

    int partition(int left, int right, vector<vector<int>> &nums){
        int pivot = left;
        while(left <= right){
            while(left <= right && nums[left][2] <= nums[pivot][2])
                left++;
            while(left <= right && nums[right][2] > nums[pivot][2])
                right--;

            if(left < right)
                swap(nums[left], nums[right]);
            else
                swap(nums[pivot], nums[right]);
        }
        return right;
    }
};