973. 最接近原点的 K 个点

题目描述

输入数组 points，points[i] = [xi, yi] 是平面上的一个点；另给一个整数 k。

点到原点 (0,0) 的距离按欧几里得距离计算，即 sqrt(x^2 + y^2)。

返回距离原点最近的 k 个点，点的顺序任意。题目保证这 k 个点组成的集合是唯一的。

示例 1

输入：points = [[1,3],[-2,2]], k = 1
输出：[[-2,2]]

解释：(-2,2) 的距离是 sqrt(8)，小于 (1,3) 的 sqrt(10)。

示例 2

输入：points = [[3,3],[5,-1],[-2,4]], k = 2
输出：[[-2,4],[3,3]]

解释：前两近的点是 (3,3) 和 (-2,4)，按任意顺序返回都可以。

数据范围

1 <= k <= points.length <= 10^4
-10^4 < xi, yi < 10^4

解题思路：大顶堆

维护一个大小不超过 k 的大顶堆，堆顶是当前候选中最远的点。逐个把点放入堆，一旦超过 k 个，就把最远的那个弹出。这样堆里始终是已扫描点中最近的 k 个。比较距离时用平方 x*x+y*y，避免开根号。Python 的 heapq 是小顶堆，所以存 -d 来模拟大顶堆。

为什么正确

处理完前 i 个点后，堆中恰好是这 i 个点里距离最小的 min(i, k) 个，堆顶是其中最远的。

复杂度分析

时间 O(n log k)，空间 O(k)。

易错点

用小顶堆存距离（不取负）会弹出最近的点，结果恰好相反；另外比较的是距离平方，不需要也不应该开根号引入浮点误差。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap = []
        for x, y in points:
            d = x * x + y * y
            heapq.heappush(heap, (-d, x, y))
            if len(heap) > k:
                heapq.heappop(heap)
        return [[x, y] for _, x, y in heap]

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        for (int[] p : points) {
            int x = p[0], y = p[1];
            int d = x * x + y * y;
            heap.offer(new int[]{d, x, y});
            if (heap.size() > k) {
                heap.poll();
            }
        }
        int[][] ans = new int[heap.size()][];
        int i = 0;
        for (int[] e : heap) {
            ans[i++] = new int[]{e[1], e[2]};
        }
        return ans;
    }
}

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        priority_queue<tuple<int, int, int>> heap;
        for (auto& p : points) {
            int x = p[0], y = p[1];
            int d = x * x + y * y;
            heap.push({d, x, y});
            if ((int)heap.size() > k) {
                heap.pop();
            }
        }
        vector<vector<int>> ans;
        while (!heap.empty()) {
            auto [d, x, y] = heap.top();
            heap.pop();
            ans.push_back({x, y});
        }
        return ans;
    }
};

type maxHeap [][3]int

func (h maxHeap) Len() int           { return len(h) }
func (h maxHeap) Less(i, j int) bool { return h[i][0] > h[j][0] }
func (h maxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *maxHeap) Push(v any)        { *h = append(*h, v.([3]int)) }
func (h *maxHeap) Pop() any          { old := *h; v := old[len(old)-1]; *h = old[:len(old)-1]; return v }

func kClosest(points [][]int, k int) [][]int {
	h := &maxHeap{}
	for _, p := range points {
		x, y := p[0], p[1]
		d := x*x + y*y
		heap.Push(h, [3]int{d, x, y})
		if h.Len() > k {
			heap.Pop(h)
		}
	}
	ans := [][]int{}
	for _, e := range *h {
		ans = append(ans, []int{e[1], e[2]})
	}
	return ans
}

const heappush = (h, e) => {
    h.push(e);
    for (let i = h.length - 1, p; i > 0 && h[p = (i - 1) >> 1][0] < h[i][0]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = h => {
    const top = h[0], last = h.pop();
    if (h.length) {
        h[0] = last;
        for (let i = 0; ;) {
            const l = 2 * i + 1, r = l + 1; let s = i;
            if (l < h.length && h[l][0] > h[s][0]) s = l;
            if (r < h.length && h[r][0] > h[s][0]) s = r;
            if (s === i) break;
            [h[s], h[i]] = [h[i], h[s]]; i = s;
        }
    }
    return top;
};

var kClosest = function(points, k) {
    const heap = [];
    for (const [x, y] of points) {
        const d = x * x + y * y;
        heappush(heap, [d, x, y]);
        if (heap.length > k) {
            heappop(heap);
        }
    }
    return heap.map(([_, x, y]) => [x, y]);
};

const heappush = (h: number[][], e: number[]) => {
    h.push(e);
    for (let i = h.length - 1, p; i > 0 && h[p = (i - 1) >> 1][0] < h[i][0]; i = p) [h[p], h[i]] = [h[i], h[p]];
};
const heappop = (h: number[][]) => {
    const top = h[0], last = h.pop()!;
    if (h.length) {
        h[0] = last;
        for (let i = 0; ;) {
            const l = 2 * i + 1, r = l + 1; let s = i;
            if (l < h.length && h[l][0] > h[s][0]) s = l;
            if (r < h.length && h[r][0] > h[s][0]) s = r;
            if (s === i) break;
            [h[s], h[i]] = [h[i], h[s]]; i = s;
        }
    }
    return top;
};

function kClosest(points: number[][], k: number): number[][] {
    const heap: number[][] = [];
    for (const [x, y] of points) {
        const d = x * x + y * y;
        heappush(heap, [d, x, y]);
        if (heap.length > k) {
            heappop(heap);
        }
    }
    return heap.map(([_, x, y]) => [x, y]);
}

载入编辑器，自己跑一遍 →

其他解法

排序最简单；大顶堆 O(n log k) 适合 k 远小于 n 或数据流式到达的情况；快速选择平均 O(n) 最快，但最坏情况较差。

解法 2：排序时间 O(n log n)，空间 O(log n)

按到原点距离的平方 x² + y² 从小到大排序，取前 k 个点。比较时使用平方而不开根号，避免浮点误差。

代码最短，思路最直接。需要对全部 n 个点排序，k 很小时比堆慢。

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        points.sort(key=lambda p: p[0] * p[0] + p[1] * p[1])
        return points[:k]

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        Arrays.sort(points, (a, b) -> Integer.compare(a[0] * a[0] + a[1] * a[1], b[0] * b[0] + b[1] * b[1]));
        return Arrays.copyOfRange(points, 0, k);
    }
}

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        sort(points.begin(), points.end(), [](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];
        });
        return vector<vector<int>>(points.begin(), points.begin() + k);
    }
};

func kClosest(points [][]int, k int) [][]int {
	sort.Slice(points, func(a, b int) bool {
		p, q := points[a], points[b]
		return p[0]*p[0]+p[1]*p[1] < q[0]*q[0]+q[1]*q[1]
	})
	return points[:k]
}

var kClosest = function(points, k) {
    points.sort((a, b) => (a[0] * a[0] + a[1] * a[1]) - (b[0] * b[0] + b[1] * b[1]));
    return points.slice(0, k);
};

function kClosest(points: number[][], k: number): number[][] {
    points.sort((a, b) => (a[0] * a[0] + a[1] * a[1]) - (b[0] * b[0] + b[1] * b[1]));
    return points.slice(0, k);
}

解法 3：快速选择时间平均 O(n)，最坏 O(n²)，空间 O(1)

用快速排序的划分思想，按距离把数组三路划分为小于、等于、大于基准的三段。若第 k 个位置落在小于段，就只在左段继续；落在大于段就只在右段继续；落在等于段则前 k 个已经是答案。最后返回前 k 个点。

平均线性时间且原地完成，是理论上最快的方法。最坏情况退化到 O(n²)，而且会打乱输入数组。

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dist(p):
            return p[0] * p[0] + p[1] * p[1]

        lo, hi = 0, len(points) - 1
        while lo < hi:
            pivot = dist(points[(lo + hi) // 2])
            lt, i, gt = lo, lo, hi
            while i <= gt:
                d = dist(points[i])
                if d < pivot:
                    points[lt], points[i] = points[i], points[lt]
                    lt += 1
                    i += 1
                elif d > pivot:
                    points[gt], points[i] = points[i], points[gt]
                    gt -= 1
                else:
                    i += 1
            if k - 1 < lt:
                hi = lt - 1
            elif k - 1 > gt:
                lo = gt + 1
            else:
                break
        return points[:k]

class Solution {
    int dist(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }

    void swap(int[][] a, int i, int j) {
        int[] t = a[i];
        a[i] = a[j];
        a[j] = t;
    }

    public int[][] kClosest(int[][] points, int k) {
        int lo = 0, hi = points.length - 1;
        while (lo < hi) {
            int pivot = dist(points[(lo + hi) / 2]);
            int lt = lo, i = lo, gt = hi;
            while (i <= gt) {
                int d = dist(points[i]);
                if (d < pivot) swap(points, lt++, i++);
                else if (d > pivot) swap(points, gt--, i);
                else i++;
            }
            if (k - 1 < lt) hi = lt - 1;
            else if (k - 1 > gt) lo = gt + 1;
            else break;
        }
        return Arrays.copyOfRange(points, 0, k);
    }
}

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        auto dist = [](const vector<int>& p) { return p[0] * p[0] + p[1] * p[1]; };
        int lo = 0, hi = points.size() - 1;
        while (lo < hi) {
            int pivot = dist(points[(lo + hi) / 2]);
            int lt = lo, i = lo, gt = hi;
            while (i <= gt) {
                int d = dist(points[i]);
                if (d < pivot) swap(points[lt++], points[i++]);
                else if (d > pivot) swap(points[gt--], points[i]);
                else i++;
            }
            if (k - 1 < lt) hi = lt - 1;
            else if (k - 1 > gt) lo = gt + 1;
            else break;
        }
        return vector<vector<int>>(points.begin(), points.begin() + k);
    }
};

func kClosest(points [][]int, k int) [][]int {
	dist := func(p []int) int { return p[0]*p[0] + p[1]*p[1] }
	lo, hi := 0, len(points)-1
	for lo < hi {
		pivot := dist(points[(lo+hi)/2])
		lt, i, gt := lo, lo, hi
		for i <= gt {
			d := dist(points[i])
			if d < pivot {
				points[lt], points[i] = points[i], points[lt]
				lt++
				i++
			} else if d > pivot {
				points[gt], points[i] = points[i], points[gt]
				gt--
			} else {
				i++
			}
		}
		if k-1 < lt {
			hi = lt - 1
		} else if k-1 > gt {
			lo = gt + 1
		} else {
			break
		}
	}
	return points[:k]
}

var kClosest = function(points, k) {
    const dist = p => p[0] * p[0] + p[1] * p[1];
    let lo = 0, hi = points.length - 1;
    while (lo < hi) {
        const pivot = dist(points[(lo + hi) >> 1]);
        let lt = lo, i = lo, gt = hi;
        while (i <= gt) {
            const d = dist(points[i]);
            if (d < pivot) {
                [points[lt], points[i]] = [points[i], points[lt]];
                lt++;
                i++;
            } else if (d > pivot) {
                [points[gt], points[i]] = [points[i], points[gt]];
                gt--;
            } else {
                i++;
            }
        }
        if (k - 1 < lt) hi = lt - 1;
        else if (k - 1 > gt) lo = gt + 1;
        else break;
    }
    return points.slice(0, k);
};

function kClosest(points: number[][], k: number): number[][] {
    const dist = (p: number[]) => p[0] * p[0] + p[1] * p[1];
    let lo = 0, hi = points.length - 1;
    while (lo < hi) {
        const pivot = dist(points[(lo + hi) >> 1]);
        let lt = lo, i = lo, gt = hi;
        while (i <= gt) {
            const d = dist(points[i]);
            if (d < pivot) {
                [points[lt], points[i]] = [points[i], points[lt]];
                lt++;
                i++;
            } else if (d > pivot) {
                [points[gt], points[i]] = [points[i], points[gt]];
                gt--;
            } else {
                i++;
            }
        }
        if (k - 1 < lt) hi = lt - 1;
        else if (k - 1 > gt) lo = gt + 1;
        else break;
    }
    return points.slice(0, k);
}

栈与堆专题的其他题目

查看栈与堆全部题目 →