K Closest Points to Origin
平面上に n 個の点があります。点 i は (xs[i], ys[i]) です。原点 (0, 0) に最も近い k 個の点を、それぞれペア [x, y] として返してください。
答えは原点からの距離が近い順に並べてください。同じ距離の点は、まず x が小さい順、次に y が小さい順にします。複数回リストに記載された点は、記載された回数分だけ数えます。
関数
- xsinteger-array
- x 座標と各点
- ysinteger-array
- 各点の y 座標
- kinteger
- 返すポイント数
- 戻り値integer-2d-array
- 最も近い k 個の点を [x, y] のペアとして、近い順に返します。同距離の場合は x、次に y の順で決まります。
制約
n == xs.length == ys.length1 ≤ k ≤ n ≤ 5000-104 ≤ xs[i], ys[i] ≤ 104- 点は重複する場合があります。重複した各点は、それぞれ別の点として扱います。
例
- 入力
- xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
- 出力
- [[-1, 2], [2, -2]]
- 説明
- 距離の二乗は10、5、8、25です。小さいほうの2つは、順に5の
(-1, 2)と8の(2, -2)です。
- 入力
- xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
- 出力
- [[-1, 1], [1, -1], [1, 1]]
- 説明
(1, 1)、(-1, 1)、(1, -1)はいずれも距離の二乗が2で、(0, 2)は4です。同点の点の中では(-1, 1)のxが最も小さく、(1, -1)はyの順で(1, 1)より前に来ます。
- 入力
- xs = [7]ys = [-7]k = 1
- 出力
- [[7, -7]]
- 説明
- 点は1つだけで、必要なのも1つなので、答えは
(7, -7)そのものです。
発展問題
点は今では1つずつ届き、そのたびに、これまでの中で最も近いk個を報告しなければなりません。どれくらいのメモリが必要ですか？
ヒント
1つずつ開いてください。開くたびに少しずつ答えに近づきます。
実際の距離を求める必要はありません。
√(x² + y²)の代わりに何を比較すれば、同じ順序が得られますか？これまでに見た中で最も近い
k個の点を保持するとします。新しい点は、そのうちのどの点と競合するでしょうか？k個以下の点を(x² + y², x, y)の順に並べた最大ヒープを保持します。新しい点がヒープの先頭の点より前に来る場合は先頭を置き換え、最後にヒープから点を取り出して、答えを後ろから埋めます。
解説
距離の計算に平方根は不要です。x² + y² は同じ順序を保ち、整数のままです。必要なのは n 個の点のうち k 個だけなので、これまでに見つかった最も近い k 個を最大ヒープに格納し、最も遠い点を先頭に置いてすぐ置き換えられるようにすれば、全件をソートする代わりに O(n log k) で済みます。
最も近い点を k 回選ぶ
正しいが、最大のテストでは終わらない
考え方
答えの最初の点はすべての点の中で最も近く、2番目の点は残りの中で最も近く、以下同様です。そこで、まだ選ばれていない点の中から (x² + y², x, y) の順で最初に来る点を毎回選び、k 回点を調べます。
各回の走査では、n 個すべての点を調べます。n = 5000、k = 2500 の場合、3要素の比較は 1.25 × 10^7 回になります。
アルゴリズム
- すべての点を未取得としてマークします。
k回繰り返します。未取得の点を調べ、(x² + y², x, y)の順で最初に来る点を見つけます。- その点を取得済みとしてマークし、答えに
[x, y]を追加します。 - 答えを返します。
def kClosest(xs, ys, k):
n = len(xs)
taken = [False] * n
result = []
for _ in range(k):
best = -1
for i in range(n):
# Compare (squared distance, x, y): the order the answer uses.
if not taken[i] and (best == -1 or
(xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
(xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
best = i
taken[best] = True
result.append([xs[best], ys[best]])
return resultすべての点を並べ替える
考え方
答えは完全な順序の先頭 k 個の点なので、すべての点を (x² + y², x, y) の順に並べ、先頭の k 個を選びます。例1では、並べ替えたキーは (5, -1, 2)、(8, 2, -2)、(10, 3, 1)、(25, 0, 5) で、先頭の2つが答えになります。
ソートの計算量は、k の値にかかわらず O(n log n) です。k が小さい場合、その処理の大部分は答えに含まれない点の順序付けに費やされます。
アルゴリズム
- 各点のキー
(x² + y², x, y)を作成します。 - キーの順に点を並べ替えます。
- 先頭から
k個を[x, y]のペアとして返します。
def kClosest(xs, ys, k):
points = list(zip(xs, ys))
# Sort by squared distance, then x, then y: the order the answer uses.
points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
return [[x, y] for x, y in points[:k]]最大ヒープ内でk個の最も近い要素を保持する
考え方
これまでに見た点のうち、最も条件に合う k 個を保持します。新しい点が押し出せるのは、保持されている点の中で順位が最も低い点だけなので、保持する点を (x² + y², x, y) の最大ヒープに格納し、その点が先頭に来るようにします。k 個の点を格納した後は、先頭の点より順位が上の点であれば先頭と置き換えます。それ以外の点は破棄します。すでに格納されている k 個の点は、その点より順位が上だからです。
例 2 では、k = 3 のとき、最初の 3 点を処理した後の先頭は (0, 2) で、キーは (4, 0, 2) です。点 (1, -1) のキーは (2, 1, -1) なので順位が上であり、先頭と置き換わります。
ヒープは部分的にしかソートされていないため、先頭を k 回取り出し、答えを最後のスロットから最初のスロットへ順に埋めます。各点で必要なヒープ操作は最大でも O(log k) なので、時間計算量は O(n log k)、メモリ計算量は O(k) です。
アルゴリズム
(x² + y², x, y)をキーとする空の最大ヒープを作ります。- 最初の
k個の点をそれぞれ挿入します。 - それ以降の各点について、トップと比較します。順位が上なら入れ替えて下方にふるい落とし、そうでなければスキップします。
- トップを
k回取り除きます。取り除いた各点を、答えの最後の空きスロットに入れます。 - 答えを返します。
import heapq
def kClosest(xs, ys, k):
# heapq is a min-heap; negating every part of (distance, x, y) puts the
# farthest kept point, in the answer's order, on top.
heap = []
for x, y in zip(xs, ys):
item = (-(x * x + y * y), -x, -y)
if len(heap) < k:
heapq.heappush(heap, item)
elif item > heap[0]:
# The new point comes before the farthest kept one: swap them.
heapq.heapreplace(heap, item)
# Popping gives the farthest first, so fill the answer from the back.
result = [None] * len(heap)
for i in range(len(result) - 1, -1, -1):
_, x, y = heapq.heappop(heap)
result[i] = [-x, -y]
return result
落とし穴と境界ケース
- 同点の場合の順序を無視する。例2には距離が2の点が3つあるため、
x、続いてyを比較しないと、入力順に並び、(-1, 1)、(1, -1)、(1, 1)の順にならない。 k個の点を格納する最小ヒープを使う。ヒープの先頭は最も近い点なので、要素を取り除くたびに答えに必要な点を捨てることになり、最後には最も遠いk個が残る。- ヒープの配列から答えをそのまま読み取る。この配列は部分的にしか順序付けられていないため、まず点を1つずつ取り除く。
- 重複する点を取り除く。
(5, 5)が3回リストされ、k = 3の場合、答えはその点が3回となる。
よくある質問4
原点に最も近い K 個の点を求めるアルゴリズムの時間計算量はどれくらいですか？
すべての点をソートするには O(n log n) かかります。ヒープ版では、走査に O(n log k)、残った k 個を並べ替えるのに O(k log k) かかり、メモリには O(k) 個の点だけを保持します。
k 個の最近傍点を求めるのに、なぜ最小ヒープではなく最大ヒープを使うのですか？
新しい点は、保持されている k 個の点のうち最も遠い点と競合するため、その点がヒープの先頭になっている必要があります。最小ヒープを使うには、n 個すべての点を保持して k 個の最も近い点を削除できるようにする必要がありますが、その場合、メモリコストは O(k) ではなく O(n) になります。
距離を比較するのに平方根が必要ですか？
いいえ。負でない数では、a < bであるのは√a < √bのときに限るため、x² + y²は距離と同じ順序になります。また、整数のままです。この場合は最大でも2 × 10^8で、32ビットの範囲内です。
原点に最も近い K 個の点は O(n) 時間で解けますか？
平均的には、はい。Quickselect は同じ (x² + y², x, y) の順序でランダムなピボットを基準に点を分割し、位置 k を含む側だけを平均 O(n) で処理し続けます。その位置より左にある k 個の点が答えです。それらをソートすると O(k log k) が追加されます。
似た問題
同じ考え方を使う問題です。2〜3問解くとパターンが身につきます。
def kClosest(xs, ys, k):
# ここにコードを書いてくださいテスト #1
テスト #2
テスト #3
引数
xs = [3, -1, 2, 0] ys = [1, 2, -2, 5] k = 2
期待される出力
[[-1, 2], [2, -2]]