# Top-k 問題

!!! question

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

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

## 方法一：走査による選択

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

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

![走査によって最大の k 個の要素を探す](top_k.assets/top_k_traversal.png)

!!! tip

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

## 方法二：ソート

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

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

![ソートによって最大の k 個の要素を探す](top_k.assets/top_k_sorting.png)

## 方法三：ヒープ

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

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

=== "<1>"
    ![ヒープに基づいて最大の k 個の要素を探す](top_k.assets/top_k_heap_step1.png)

=== "<2>"
    ![top_k_heap_step2](top_k.assets/top_k_heap_step2.png)

=== "<3>"
    ![top_k_heap_step3](top_k.assets/top_k_heap_step3.png)

=== "<4>"
    ![top_k_heap_step4](top_k.assets/top_k_heap_step4.png)

=== "<5>"
    ![top_k_heap_step5](top_k.assets/top_k_heap_step5.png)

=== "<6>"
    ![top_k_heap_step6](top_k.assets/top_k_heap_step6.png)

=== "<7>"
    ![top_k_heap_step7](top_k.assets/top_k_heap_step7.png)

=== "<8>"
    ![top_k_heap_step8](top_k.assets/top_k_heap_step8.png)

=== "<9>"
    ![top_k_heap_step9](top_k.assets/top_k_heap_step9.png)

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

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

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

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