ヒープで原点にいちばん近い k 個の点を選ぶ

ヒープ問題 110

平面の点の並びから、原点 (0, 0) に近い方の k 個を選んで返します。全部を並べ替えず、近い k 個だけを「いちばん遠い物が根に来るヒープ」に持ち続ける練習です。

海外で定番の問題リスト 4 本には入っていません

問題

配達の拠点が原点 (0, 0) にあり、届け先の場所が平面の点 [x, y] の並び points で渡されます。拠点からまっすぐ測って近い方から k 件を選び、その点を並びにして返す関数を書いてください。

近さは、まっすぐの距離 (x の 2 乗と y の 2 乗の和の平方根) で比べます。平方根を取らず、2 乗の和のまま比べてかまいません
返す順は、近い順です。距離が同じ点どうしは、points で前にある方を先にします。選ぶときも同じ決まりで、距離が同じなら前にある点を選びます
points の長さは 1 以上 10000 以下、x と y はどれも -10000 以上 10000 以下の整数です
k は 1 以上、points の長さ以下です
全部を並べ替える O(n log n) ではなく、ヒープを使って O(n log k) で選んでください

入力と出力の例

例 1

2 乗の和は順に 10、4、25、2 です。近い 2 件は、2 の [-1, 1] と 4 の [2, 0] です。

points = [[3, 1], [2, 0], [0, -5], [-1, 1]] k = 2

[[-1, 1], [2, 0]]

例 2

[0, 2] と [2, 0] と [-2, 0] は、どれも 2 乗の和が 4 で同じです。前にある 2 つを選びます。

points = [[0, 2], [2, 0], [-2, 0]] k = 2

[[0, 2], [2, 0]]

例 3

点が 1 つで、k も 1 です。

points = [[4, -3]] k = 1

[[4, -3]]

面接での聞かれ方

面接では、短く次のように渡されることが多いです。

「平面の点の並びから、原点に近い k 個を返してください。並べ替えずにできますか？」
ヒープで書くと「なぜ近い順ではなく、遠い順のヒープなのですか？」と理由を聞かれます
「平方根は要りますか？」「点が何億個も流れてくるとしたら？」と広げて聞かれます

返す順と、距離が同じ点の扱いを最初に確かめておくと、答えが 1 つに決まり、例で比べやすくなります。

ヒント

まだ開いていません

付箋は 4 枚。下の段ほど答えに近づきます

ヒント 2
ヒント 3
ヒント 4答えに近い

アルゴリズムのチートシートの「スタックとキュー」に、この型の要点を表とコードでまとめています。

スタックは後から入れた物を先に、キューは先に入れた物を先に取り出します。括弧の対応や、順番待ちの処理の形です。

まず考える手順

段 1

まずは全部の点を距離で並べ替え、前から k 個を取る書き方です。短く正しく書けますが O(n log n) で、要らない遠い点の順まで決めています。

段 2

「今までに見た点のうち、近い方の k 個」だけを持ち続けます。新しい点が来たら、持っている k 個のうちいちばん遠い点と比べ、それより近ければ入れ替えます。いちばん遠い点をすぐ見たいので、遠い順のヒープ (根にいちばん遠い点がある木) を使います。1 回の入れ替えは O(log k) です。

段 3

距離が同じ点の決まりも、比べ方に入れます。「遠い」を「距離が大きい、距離が同じなら points で後ろにある」と決めると、根にはいつも「いちばん外したい点」が来ます。新しい点は、根より距離が小さいときだけ入れ替えます。新しい点は根より後ろにあるので、距離が同じなら入れ替えません。

段 4

最後にヒープに残った k 個を、距離、points の中の位置の順で並べ替えて返します。この並べ替えは k 個だけなので O(k log k) です。ヒープには点そのものではなく、points の中の番号を入れておくと、距離が同じときの比べ方にそのまま使えます。

解答

Python

from heapq import heappush
from heapq import heapreplace


def k_closest(points, k):
    h = []
    for i, (x, y) in enumerate(
            points):
        d = x * x + y * y
        # 遠くて後ろの点が根に来る
        key = (-d, -i)
        if len(h) < k:
            heappush(h, key)
        elif d < -h[0][0]:
            heapreplace(h, key)
    h.sort(reverse=True)
    return [points[-i]
            for _, i in h]

JavaScript

function kClosest(points, k) {
  const d = points.map(
    ([x, y]) => x * x + y * y,
  );
  // a が b より遠い (外したい)
  const far = (a, b) =>
    d[a] > d[b] ||
    (d[a] === d[b] && a > b);
  const h = [];
  for (const i of d.keys()) {
    if (h.length < k) {
      h.push(i);
      up(h, h.length - 1, far);
    } else if (d[i] < d[h[0]]) {
      h[0] = i;
      down(h, 0, far);
    }
  }
  h.sort((a, b) =>
    far(a, b) ? 1 : -1);
  return h.map(
    (i) => points[i]);
}

function up(h, i, far) {
  while (i > 0) {
    const p = (i - 1) >> 1;
    if (!far(h[i], h[p])) {
      return;
    }
    [h[p], h[i]] = [h[i], h[p]];
    i = p;
  }
}

function down(h, i, far) {
  const n = h.length;
  while (true) {
    let s = i;
    const l = 2 * i + 1;
    const r = l + 1;
    if (l < n &&
        far(h[l], h[s])) {
      s = l;
    }
    if (r < n &&
        far(h[r], h[s])) {
      s = r;
    }
    if (s === i) {
      return;
    }
    [h[s], h[i]] = [h[i], h[s]];
    i = s;
  }
}

Go

import (
	"container/heap"
	"sort"
)

type farHeap struct {
	ids  []int
	dist []int
}

func (h *farHeap) Len() int {
	return len(h.ids)
}

func (h *farHeap) Less(
	i, j int,
) bool {
	a, b := h.ids[i], h.ids[j]
	if h.dist[a] != h.dist[b] {
		return h.dist[a] > h.dist[b]
	}
	return a > b
}

func (h *farHeap) Swap(
	i, j int,
) {
	s := h.ids
	s[i], s[j] = s[j], s[i]
}

func (h *farHeap) Push(x any) {
	h.ids = append(h.ids, x.(int))
}

func (h *farHeap) Pop() any {
	n := len(h.ids) - 1
	x := h.ids[n]
	h.ids = h.ids[:n]
	return x
}

func kClosest(
	points [][]int, k int,
) [][]int {
	d := make([]int, len(points))
	for i, p := range points {
		d[i] = p[0]*p[0] + p[1]*p[1]
	}
	h := &farHeap{dist: d}
	for i := range points {
		if h.Len() < k {
			heap.Push(h, i)
			continue
		}
		if d[i] < d[h.ids[0]] {
			h.ids[0] = i
			heap.Fix(h, 0)
		}
	}
	ids := h.ids
	sort.Slice(ids, func(
		a, b int,
	) bool {
		return h.Less(b, a)
	})
	out := [][]int{}
	for _, i := range ids {
		out = append(out, points[i])
	}
	return out
}

Python の heapq は小さい順なので、距離と番号の両方に - を付けて、遠くて後ろの点が根に来るようにしています。Go の最後の並べ替えは、ヒープの Less (遠い方が先) を逆に使って、近い順にしています。

計算量

時間は O(n log k) です。n 個の点それぞれについて、大きさ k 以下のヒープを直すのが O(log k) です。最後の並べ替えは O(k log k) で、k は n 以下なのでこれに含まれます
使う記憶は O(k) です。ヒープに持つのは k 個の番号までです (JavaScript と Go は距離の並びに O(n) を使いますが、距離をその場で計算すれば O(k) にできます)

よくある間違い

近い順のヒープに全部入れる。全部入れてから k 回取り出すと、記憶が O(n) になります。動きはしますが、遠い順のヒープに k 個だけ持つ書き方との違いを言えるようにします
平方根を取る。 sqrt は小数になり、同じ距離が丸めでずれることがあります。2 乗の和のまま比べれば整数で済みます
距離が同じ点の扱いを決めない。例 2 で [-2, 0] を返す書き方も「近い 2 件」としては正しく見えますが、この問題の決まりでは外れます。比べ方に番号を入れて、1 つに決めます
最後に並べ替えない。ヒープの中は近い順に並んではいません。返す前に並べ替えます

解答を動かした版 Python 3.14.7Node.js 24.22.1Go 1.24.2確認日 2026年11月8日

main.py

def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    """
    Goal: 何を返すか

    idea1: まず思いつく案
    -> runtime: O(?)  space: O(?)

    idea2: もっと速い案
    -> runtime: O(?)  space: O(?)

    edge case: 空、1 つだけ、など
    """
    # ここに書きます
    pass


# 「実行」を押すと、この行の結果が出ます
print(k_closest([[3, 1], [2, 0], [0, -5], [-1, 1]], 2))

できているか「例で確かめる」を押すと、問題の例を全部この関数に渡して、通ったかをここに出します