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 問題をより効率的に解くことができます。手順は以下の図のとおりです。

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

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

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

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

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