Kth Largest Element in an Array
整数配列 nums と整数 k が与えられます。nums の中で k 番目に大きい値を返してください。つまり、配列を大きい順に並べたとき、先頭から数えて k 番目の値です。
同じ値もそれぞれ別個に数えます。[5, 5, 1] では、最大の値は 5 で、2 番目に大きい値も 5 です。
関数
- numsinteger-array
- ランク付けする値
- kinteger
- 返す最大値。最大値の場合は1
- 戻り値integer
- 重複を数えた場合の、k 番目に大きい値
制約
1 ≤ k ≤ nums.length ≤ 104-104 ≤ nums[i] ≤ 104- 等しい値も、それぞれ別の値として数えます。
例
- 入力
- nums = [7, 2, 9, 4, 9, 1]k = 2
- 出力
- 9
- 説明
- 値を大きい順に並べると、
9, 9, 7, 4, 2, 1です。9は2つそれぞれ別に数えるので、2番目に大きい値は7ではなく9です。
- 入力
- nums = [5, -3, 8, 0, 2]k = 4
- 出力
- 0
- 説明
- 値を大きい順に並べると、
8, 5, 2, 0, -3となり、その4番目の値は0です。
- 入力
- nums = [6]k = 1
- 出力
- 6
- 説明
- 値が 1 つで、
k = 1の場合、その値が最大です。
提出時に隠しテスト+15件
発展問題
値が1つずつ届くようになりました。各値が届くたびに、それまでに見たすべての値の中央値を、値ごとに O(log n) の時間で報告できますか？
ヒント
1つずつ開いてください。開くたびに少しずつ答えに近づきます。
大きい順に並べると、答えは既知の位置にあります。それはどれでしょうか？そして、それを知るために他のすべての値が必要でしょうか？
k番目に大きい値は、
k個の最大値の中で最小の値です。これまでに見たk個の最大値だけを保持する場合、新しい値をそれらのうちどの値と比較しますか？最大でも
k個の値を持つ最小ヒープを維持します。新しい値がヒープの先頭より大きい場合は置き換え、最後に先頭にある値が答えです。平均時間計算量をO(n)にするには、quicksortと同様にランダムなピボットを基準に分割し、インデックスn-kを含む側だけを残します。
解説
ソートして1つの位置を読む方法で答えが得られ、ここでは十分に高速です。面接官が見たいのは、どれだけ並べ替えを省略できるかです。必要なのは1つの位置であって、すべてのnではないからです。サイズkの最小ヒープは、答えになり得る値だけを保持します。また、クイックセレクトはクイックソートのように分割しますが、答えを含む側だけをたどるため、平均実行時間をO(n)まで短縮できます。
1つの位置で並べ替えて読み取る
考え方
k 番目に大きい値は、ソート順で定義されるため、その順序を作ります。大きい順にソートすると、[7, 2, 9, 4, 9, 1] は [9, 9, 7, 4, 2, 1] になり、k 番目に大きい値はインデックス k-1 にあります。k = 2 の場合はインデックス 1、2 番目の 9 です。小さい順にソートする場合は、代わりにインデックス n-k を読み取ります。[1, 2, 4, 7, 9, 9] のインデックス 4 にある値も、同じ 9 です。
重複値に特別な処理は必要ありません。ソートでは各値がすべて保持され、それぞれが独自の位置を占めます。
n = 10^4 の場合、ソートに必要な比較回数は約 n log n ≈ 1.3 × 10^5 回で、すべてのテストに合格します。無駄なのは、必要なのは 1 つの位置だけなのに、n 個すべての値を順序付けることです。次の 2 つの方法では、その作業を減らします。
アルゴリズム
numsをコピーして、呼び出し元の配列を元のままにします。- コピーをソートします。数値比較を使ってください。言語によっては、デフォルトで数値をテキストとして比較するものがあります。
- 大きい順の場合はインデックス
k-1、小さい順の場合はインデックスn-kを返します。
def findKthLargest(nums, k):
# Largest first: the k-th largest sits at index k-1.
ordered = sorted(nums, reverse=True)
return ordered[k - 1]最小ヒープで上位 k 個を保持する
考え方
k番目に大きい値は、k個の最大値の中で最も小さい値です。そこで、numsを一度だけ走査し、これまでに見た値のうち大きい方からk個だけを最小ヒープに保持します。最小ヒープの先頭は最小値なので、これが答えの候補になります。
値xが到着したとき、ヒープに保持されている値がk個未満なら、xを追加します。それ以外の場合は、xと先頭の値を比較します。xがそれより大きくなければ、保持している少なくともk個の値がx以上なので、xが答えになることはなく、スキップします。xの方が大きければ、先頭の値は大きい方からk個の範囲から外れるため、xで置き換えます。例2ではk = 4なので、最初の4つの値5, -3, 8, 0でヒープが埋まり、先頭は-3です。次に2は-3より大きいためこれと置き換わり、先頭は0になります。そして0が答えです。
各値にかかるヒープ操作は最大1回で、その計算量はO(log k)です。したがって、全体では時間計算量がO(n log k)、空間計算量がO(k)です。kが小さい場合はソートより効率的で、ストリームにも対応できます。すべての値を一度に保持する必要はありません。Pythonにはheapq、JavaにはPriorityQueue、C++にはgreaterと組み合わせるpriority_queue、Goにはcontainer/heap、RustにはReverseと組み合わせるBinaryHeap、PHPにはSplMinHeapがあります。他の言語のコードでは、配列を使ってヒープを実装しています。インデックスiの子は2i+1と2i+2にあります。1から数えるLuaとRでは、2iと2i+1にあります。
アルゴリズム
- 空の最小ヒープから始めます。
- 各値
xについて、ヒープにk個未満の値しかない間は、その値を追加します。 k個になったら、xがヒープの先頭より大きい場合にのみ、先頭をxに置き換えます。- 最後の値を処理した後、ヒープの先頭を返します。
import heapq
def findKthLargest(nums, k):
# A min-heap of the k largest values so far; its top is the smallest of them.
heap = []
for x in nums:
if len(heap) < k:
heapq.heappush(heap, x)
elif x > heap[0]:
heapq.heapreplace(heap, x) # drop the top, add x
return heap[0]3分割パーティションを用いたクイックセレクト
考え方
クイックソートはピボットを選び、分割します。小さい値は左側に、大きい値は右側に置きます。1回分割すると、どちら側もまだソートされていなくても、ピボットは最終的なソート後のインデックスに位置します。クイックセレクトはこの性質を利用します。小さい順では、答えはインデックス target = n-k にあります。分割後、target はピボットの左、ピボット位置、またはピボットの右のいずれかにあるため、片側だけを続け、もう片側を切り捨てます。
[7, 2, 9, 4, 9, 1] と k = 2 の場合、target は 6-2 = 4 です。4 を基準に分割すると、2 と 1 はインデックス 0 と 1 に、4 はインデックス 2 に、7, 9, 9 はインデックス 3 から 5 に配置されます。インデックス 4 は右側にあるので、インデックス 3 から 5 だけを残します。それらを 9 を基準に分割すると、7 はインデックス 3 に、2つの9はインデックス 4 と 5 に配置されます。インデックス 4 にあるのは9なので、答えは 9 です。
3方向分割を使います。ピボットより小さい値、ピボットと等しい値、ピボットより大きい値の順に並べ、lt と gt で管理します。等しい値のブロック [lt, gt] はソート済みの位置にあるため、target がその範囲内にあれば処理は完了です。単純な2方向分割では、10^4 個の 7 のコピーを含む配列は、1回ごとに値が1つしか減らず、約 5 × 10^7 ステップかかります。3方向分割なら1回の走査で答えが得られます。
ピボットはランダムに選びます。半分の確率で範囲の中央半分に入るため、範囲は最大でも4分の3まで縮小されます。そのため、期待される処理量は n 個の値を数回走査する程度、つまり O(n) です。すべてのピボットが端の値になると、最悪の場合は依然として O(n²) です。また、先頭の要素など固定の選び方では、入力がソート済みの場合にこのケースが発生します。このコードはコピーに対して動作するため、メモリを O(n) 使用します。入力を変更してよければ、nums 自体を分割することで O(1) にできます。
アルゴリズム
numsをaにコピーし、target = n-k、lo = 0、hi = n-1を設定します。a[lo..hi]からランダムなピボットを選びます。a[lo..hi]をピボットより小さい値、等しい値、大きい値に分割し、等しい値をa[lt..gt]に残します。target < ltの場合はhi = lt-1を設定します。target > gtの場合はlo = gt+1を設定します。それ以外の場合はピボットを返します。- 手順 2 から繰り返します。
import random
def findKthLargest(nums, k):
a = list(nums)
target = len(a) - k # the answer's index once a is sorted smallest first
lo, hi = 0, len(a) - 1
while True:
pivot = a[random.randint(lo, hi)]
# Three-way partition of a[lo..hi]: < pivot, then == pivot, then > pivot.
lt, i, gt = lo, lo, hi
while i <= gt:
if a[i] < pivot:
a[lt], a[i] = a[i], a[lt]
lt += 1
i += 1
elif a[i] > pivot:
a[i], a[gt] = a[gt], a[i]
gt -= 1
else:
i += 1
# Now a[lt..gt] all equal pivot, and they are in their sorted places.
if target < lt:
hi = lt - 1
elif target > gt:
lo = gt + 1
else:
return pivot
落とし穴と境界ケース
間違った答えの多くは、重複要素と、位置を数える2つの方法の混同によって生じます。
- 最初に重複要素を削除する。この問題では各要素をすべて数えます。
[7, 2, 9, 4, 9, 1]でk = 2の場合、答えは9ですが、配列をセットに変換すると7になります。 - 誤ったインデックスを読む。
kは1から数えるため、答えは大きい順ではインデックスk-1に、小さい順ではインデックスn-kにあります。n-k-1ではありません。 - 数値を文字列としてソートする。JavaScriptとTypeScriptでは、
[10, 9, 2].sort()の結果は[10, 2, 9]です。(a, b) => a - bを渡してください。 - サイズ
kの最大ヒープを使う。最大値を取り除くと、k個の最小値が残り、k番目に小さい値が返されます。 - 二分割方式の分割や固定ピボットを使ったクイックセレクト。等しい値が多数ある場合や配列がソート済みの場合、計算量が
O(n²)になり、大規模なテストにはそのようなケースが含まれています。
よくある質問4
配列内のK番目に大きい要素の時間計算量は何ですか？
ソートには O(n log n) の時間がかかります。サイズ k の最小ヒープには、O(n log k) の時間と O(k) のメモリが必要です。ランダムなピボットを使うクイックセレクトは、平均で O(n) の時間、最悪の場合は O(n²) の時間がかかりますが、ランダムなピボットを使うことで最悪の場合が発生する可能性は非常に低くなります。
k番目に大きい要素を見つけるのに、なぜ最大ヒープではなく最小ヒープを使うのですか？
ヒープには、これまでに見つかった最大の値 k 個が格納されます。その中で比較して追い出す必要があるのは最小の値です。最小ヒープでは、その値が先頭に置かれます。最大ヒープを使うには、n 個すべての値を入れて k-1 回取り出す必要があり、O(n) のメモリが必要です。
k 番目に大きい要素を求めるには、ヒープとクイックセレクトのどちらを使うべきですか？
Quickselectは平均するとより高速で、O(n)ですが、すべての値をメモリ上に保持し、それらを並べ替える必要があります。ヒープは最悪の場合でも性能が悪化せず、O(n log k)で、値が一つずつ届き、すべてを保存できない場合にも使えます。面接では両方を説明し、追加質問で求められた方をコードにしましょう。
最悪の場合でも、k 番目に大きい要素を線形時間で見つけられますか？
はい。中央値の中央値ルールでは、値の一定の割合を必ず分割できるピボットを選ぶため、最悪の場合でも選択はO(n)になりますが、実際にはランダムなピボットより遅くなります。値が-10^4から10^4までに限られている場合は、各値の出現回数を数え、k個の値を通過するまで10^4から順にたどる方法もあり、時間計算量はO(n + 2 × 10^4)です。
似た問題
同じ考え方を使う問題です。2〜3問解くとパターンが身につきます。
Python
def findKthLargest(nums, k):
# ここにコードを書いてくださいケース1
ケース2
ケース3
入力
nums = [7, 2, 9, 4, 9, 1] k = 2
期待値
9