Kth Largest Element in a Stream
整数のストリームは、numsの値から始まります。次に、addsの値が1つずつ到着します。各値が到着するたびに、同じ値も別々に数え、これまでに見たすべての値の中でk番目に大きい値を記録します。addsの各値に対して1つずつ、記録した値を順番に返します。
numsには少なくともk-1個の値が含まれているため、最初の値が到着した後は、常に選択できる値が少なくともk個あります。
関数
- kinteger
- 報告する順位。最大値の場合は1
- numsinteger-array
- ストリームにすでに含まれている値
- addsinteger-array
- 順番に届く値
- 戻り値integer-array
- 各要素の到着後の k 番目に大きい値
制約
1 ≤ nums.length ≤ 1041 ≤ adds.length ≤ 1041 ≤ k ≤ nums.length + 1-104 ≤ nums[i], adds[i] ≤ 104numsはどのような順序でもよく、同じ値もそれぞれ別の値として数えます。
例
- 入力
- k = 3nums = [7, 1, 5]adds = [3, 9, 6, 2, 8]
- 出力
- [3, 5, 6, 6, 7]
- 説明
- 3の後には、7、5、3、1の中で3番目に大きい数である3が来ます。次に9によって5に増え、6によって6に増え、2では変化せず、8によって7に増えます（9、8、7）。
- 入力
- k = 1nums = [-2]adds = [-5, 0, -1, 4]
- 出力
- [-2, 0, 0, 4]
- 説明
k = 1の場合、答えはそれまでの最大値です。0が現れるまでは-2、-1の後も0のままで、その後は4です。
- 入力
- k = 4nums = [6, 6, 2]adds = [6, 1, 9]
- 出力
- [2, 2, 6]
- 説明
- 最初の到着で4つの値、6、6、6、2ができるので、4番目に大きい値は2です。また、1を加えても変わりません。9の後では、4つの最大値は9、6、6、6です。6はすべての出現回数を数えます。
発展問題
ストリームから値を削除することもできるとしたら、同じ質問にどう答えますか？
ヒント
1つずつ開いてください。開くたびに少しずつ答えに近づきます。
到着のたびにすべてを並べ替えれば正しい答えが得られますが、そのたびに作業のほとんどをやり直しています。後になっても、どの値が
k番目に大きい値になり得るでしょうか？値は次々に追加され、取り除かれることはありません。ある値より大きな値がその上に
k個現れたら、その値が再び答えになることはありません。これまでに見た値のうち、最も大きいk個だけを保持すれば十分です。これらの
k個の値を最小ヒープに保持します。その先頭はそれらの中で最小の値であり、全体ではk番目に大きい値です。値が到着したら、ヒープに保持されている値がk個未満の間は追加し、その後は新しい値の方が大きい場合に限り先頭の値を置き換えます。
解説
値がストリームから出ていくことはないため、現在の k 番目に大きい値を下回った値は、以後候補から外れます。最小ヒープには最大の値を k 個だけ保持します。ヒープの先頭が答えとなり、新しい値が到着するたびに必要なのは比較1回と、最大1回の O(log k) のヒープ更新です。
到着するたびに並べ替える
正しいが、最大のテストでは終わらない
考え方
すべての値をリストに保持します。値が追加されるたびに昇順に並べ替え、定義上 k 番目に大きい値であるインデックス length-k を読み取ります。
値が追加されるたびに、最大で n + m 個の値を並べ替えます。ここで、n と m はそれぞれ nums と adds の長さです。それぞれが 10^4 個の場合、合計で約 2 × 10^9 回の比較になります。
アルゴリズム
numsをリストvaluesにコピーします。addsの各値をvaluesに追加します。valuesを昇順に並べ替え、values[length-k]を記録します。- 記録した値を返します。
def kthLargestStream(k, nums, adds):
values = list(nums)
result = []
for x in adds:
values.append(x)
values.sort()
# In increasing order the k-th largest sits k places from the end.
result.append(values[len(values) - k])
return resultk 個の最大値からなる最小ヒープ
考え方
ある値より大きい値が k 個現れたら、ストリームの残りの部分でもその状態は変わらないため、これまでに見た k 個の最大値だけが重要です。それらを最小ヒープに格納します。ヒープの先頭は k 個の最大値のうち最小の値であり、これは全体で k 番目に大きい値そのものです。
x が到着し、ヒープが満杯の場合は、先頭の値と比較します。x がそれより大きくなければ、k 個の最大値に含まれないため、何も変わりません。大きければ、先頭の値が外れます。先頭の値を x に置き換えて、O(log k) で下方向にふるい落とします。最初の値にも同じ処理を適用するため、nums は k より長くてもかまいません。
最初の例では、nums の処理後、ヒープには 7、1、5 が格納され、先頭は 1 です。3 が到着すると 1 と置き換わり、先頭の値 3 が最初の答えになります。
アルゴリズム
- 空の最小ヒープから始めます。
xを追加するには、ヒープに保持されている値がk個未満の場合は、xをプッシュします。そうでない場合、xがトップより大きければ、トップをxに置き換えてヒープの順序を復元します。- この方法で
numsのすべての値を追加します。 addsの各値について、それを追加し、ヒープのトップを記録します。- 記録した値を返します。
import heapq
def kthLargestStream(k, nums, adds):
# A min-heap of the k largest values so far; heap[0] is the smallest of them.
heap = []
def add(x):
if len(heap) < k:
heapq.heappush(heap, x)
elif x > heap[0]:
heapq.heapreplace(heap, x) # drop the top, add x
for x in nums:
add(x)
result = []
for x in adds:
add(x)
result.append(heap[0]) # the k-th largest so far
return result
落とし穴と境界ケース
- 重複を取り除いてしまう。9、6、6、6、2のうち大きい方から4つは、9、6、6、6です。集合を使ったり、すでに存在する値をスキップするチェックを行ったりすると、3つ目の例では6ではなく2が返されます。
- 追加する前に答えを読み取る。
numsにk-1個の値が入っている場合、ヒープが満杯になるのは最初の要素を追加した後だけなので、先に追加してから先頭の値を読み取ります。 - ソート後に反対側から数えてしまう。昇順では、
k番目に大きい値はインデックスlength-kにあります。インデックスk-1にあるのは、k番目に小さい値です。 - デフォルトのヒープ順序を使ってしまう。C++ の
priority_queueと Rust のBinaryHeapは最大ヒープです。greater<int>を渡すか、値をReverseでラップします。
よくある質問3
k 番目に大きい要素を求めるのに、なぜ最大ヒープではなく最小ヒープを使うのですか？
k 個の最大値のうち最小の値が答えであり、新たに加わる値が超える必要のある値でもあります。サイズ k の最小ヒープを使えば、その値を先頭に保てます。すべての値を格納する最大ヒープでは最大値が先頭になるため、k 番目の値にたどり着くには、値が追加されるたびに k-1 個の値を取り出して戻す必要があります。
ストリーム中のK番目に大きい要素の時間計算量はどれくらいですか？
到着のたびにO(log k)のコストがかかります。先頭との比較が1回と、k個の値を含むヒープでの sift が最大1回です。numsからヒープを構築するコストはO(n log k)なので、全体では時間計算量がO((n + m) log k)、追加の領域計算量がO(k)です。
ストリーム内のK番目に大きい要素は、ソート済みリストで解けますか？
はい。すべての値をソートした状態に保ち、各到着値を二分探索で適切な位置に挿入し、末尾からk番目の値を読み取ります。探索はO(log n)ですが、挿入位置より後ろの値をずらす処理は到着ごとにO(n)です。開始時の値が10^4個、到着値が10^4個の場合、シフト回数は約1.5 × 10^8回です。
似た問題
同じ考え方を使う問題です。2〜3問解くとパターンが身につきます。
def kthLargestStream(k, nums, adds):
# ここにコードを書いてくださいテスト #1
テスト #2
テスト #3
引数
k = 3 nums = [7, 1, 5] adds = [3, 9, 6, 2, 8]
期待される出力
[3, 5, 6, 6, 7]