hello-algo/ja/docs/chapter_heap/top_k.md

Top-k 問題

!!! question

長さ $n$ の未整列配列 `nums` が与えられたとき、配列内で最大の $k$ 個の要素を返してください。

この問題について、まずは発想が比較的直接的な 2 つの解法を紹介し、その後でより効率の高いヒープ解法を紹介します。

方法一：走査による選択

以下の図に示すように k 回の走査を行い、各ラウンドでそれぞれ第 $1$、$2$、$\dots$、k 位の要素を取り出すことができます。時間計算量は O(nk) です。

この方法は k \ll n の場合にしか適していません。k が n にかなり近いと、時間計算量は O(n^2) に近づき、非常に時間がかかるためです。

!!! tip

$k = n$ のとき、完全な昇順列を得ることができ、この場合は「選択ソート」アルゴリズムと等価になります。

方法二：ソート

以下の図に示すように、まず配列 nums をソートし、その後で右端の k 個の要素を返すことができます。時間計算量は O(n \log n) です。

明らかに、この方法は必要以上の処理を行っています。なぜなら、必要なのは最大の k 個の要素を見つけることだけであり、他の要素をソートする必要はないからです。

方法三：ヒープ

ヒープを用いることで、Top-k 問題をより効率的に解くことができます。手順は以下の図のとおりです。

1. 最小ヒープを初期化し、そのヒープ頂点の要素が最小となるようにします。
2. まず配列の先頭 k 個の要素を順にヒープへ挿入します。
3. k + 1 番目の要素から開始し、現在の要素がヒープ頂点の要素より大きければ、ヒープ頂点の要素を取り出し、現在の要素をヒープへ挿入します。
4. 走査が完了した後、ヒープに保持されているのが最大の k 個の要素です。

=== "<1>"

=== "<2>"

=== "<3>"

=== "<4>"

=== "<5>"

=== "<6>"

=== "<7>"

=== "<8>"

=== "<9>"

サンプルコードは以下のとおりです。

[file]{top_k}-[class]{}-[func]{top_k_heap}

合計で n 回のヒープ挿入と取り出しを行い、ヒープの最大長は k であるため、時間計算量は O(n \log k) です。この方法は非常に効率が高く、k が小さいときは時間計算量が O(n) に近づき、k が大きいときでも O(n \log n) を超えることはありません。

さらに、この方法は動的データストリームの利用シーンにも適しています。データが継続的に追加される場合でも、ヒープ内の要素を保ち続けることで、最大の k 個の要素を動的に更新できます。