C++でTop Kの頻出単語を抽出するアルゴリズムの解説

問題の概要

空でない単語のリストが与えられたとき、出現頻度の高い上位k個の単語を見つける問題を考えてみましょう。答えは頻度の高い順（降順）にソートされている必要があります。また、2つの単語の頻度が同じ場合は、アルファベット順で先に来る単語を先に配置します。

例えば、配列が ['the', 'sky', 'is', 'blue', 'the', 'weather', 'is', 'comfortable'] の場合、最も頻出する単語は ["is", "the", "blue"] となります。

解決手順

この問題を解くには、以下の手順に従います。

• マップ m を1つ定義する
• 優先度付きキュー v を1つ作成する
• i を 0 から n（単語配列のサイズ）まで繰り返し、m[words[i]] を1ずつ増やす
• マップ内の各要素 e について以下を処理する
  • v のサイズが k 未満であれば、e を v に挿入する
  • そうでなく、v.top() の頻度値が e の頻度値より小さい場合は、v から先頭要素を削除して e を挿入する
  • そうでなく、v.top() の頻度値が e の頻度値と等しく、かつ v.top() のキー（単語）が e のキーより大きい（アルファベット順で後）場合は、v から先頭要素を削除して e を挿入する
• 配列 res を1つ定義する
• v が空でない間、以下を繰り返す
  • temp に v の先頭要素を代入する
  • v から先頭要素を削除する
  • temp のキー（単語）を res 配列に挿入する
• res 配列を反転する
• res を返す

C++での実装例

理解を深めるために、以下の実装例を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
struct Comparator{
    bool operator()(pair <string ,int> a, pair <string, int> b){
        if(a.second != b.second) return !(a.second < b.second);
            return !(a.first > b.first);
    }
};
class Solution {
public:
    static bool cmp(pair <string, int> a, pair <string, int> b){
        if(a.second != b.second) return a.second > b.second;;
            return a.first < b.first;
    }
    vector<string> topKFrequent(vector<string>& words, int k) {
        map<string, int> m;
        priority_queue < pair <string, int>, vector < pair <string, int> >, Comparator > v;
        for(int i = 0; i < words.size(); i++){
            m[words[i]]++;
        }
        map<string, int> :: iterator i = m.begin();
        while(i != m.end()){
            if(v.size() < k){
                v.push(*i);
            }
            else if(v.top().second < i->second){
                v.pop();
                v.push(*i);
            }
            else if(v.top().second == i->second && v.top().first > i->first){
                v.pop();
                v.push(*i);
            }
            i++;
        }
        vector <string> res;
        while(!v.empty()){
            pair <string, int> temp = v.top();
            v.pop();
            res.push_back(temp.first);
        }
        reverse(res.begin(), res.end());
        return res;
    }
};
main(){
    Solution ob;
    vector<string> v = {"the", "sky", "is", "blue", "the", "weather", "is", "comfortable"};
    print_vector(ob.topKFrequent(v, 3));
}

入力

["the", "sky", "is", "blue", "the", "weather", "is", "comfortable"]
3

出力

["is","the","blue"]

アルゴリズムのポイント

この実装では、まず map を使って各単語の出現回数をカウントします。次に、カスタムコンパレータを持つ priority_queue を利用することで、頻度が低い単語やアルファベット順で後ろの単語を効率的にキューの先頭に保持できます。これにより、キューのサイズを k に制限しながら、頻度上位 k 個の単語だけを残すことが可能になります。

最後にキューから要素を取り出して配列に格納し、反転することで頻度の降順かつアルファベット昇順の結果を得られます。計算量は、単語数を n、異なる単語数を m とすると、O(n + m log k) となり、効率的な処理が実現できます。