Pythonでバイグラムの直後に現れる単語を抽出する方法

問題概要

あるテキストと、2つの単語「first」「second」が与えられたとします。このとき、テキストの中から「first second third」という形式で並んでいる箇所を探します。ここで「second」は「first」の直後に位置し、「third」は「second」の直後に位置します。

該当する箇所が見つかるたびに、その「third」にあたる単語を答えのリストへ追加し、最終的にそのリストを返します。

例えば、テキストが「lina is a good girl she is a good singer」、first =「a」、second =「good」である場合、答えは [girl, singer] となります。

解法のアプローチ

この問題は、以下の手順で解くことができます。

• テキストをスペースで分割し、単語のリストに変換する
• 結果を格納するための空のリスト res を用意する
• i を 0 からテキストの長さ - 1 までループさせる
  • i + 2 がテキストの長さ未満であり、かつ text[i] が first と一致し、text[i + 1] が second と一致する場合、text[i + 2] を res に追加する
• 最後に res を返す

実装例

理解を深めるために、以下のPythonコードを見てみましょう。

class Solution(object):
    def findOcurrences(self, text, first, second):
        text = text.split(" ")
        res = []
        for i in range(len(text)):
            if i + 2 < len(text) and text[i] == first and text[i + 1] == second:
                res.append(text[i + 2])
        return res

ob1 = Solution()
print(ob1.findOcurrences("lina is a good girl she is a good singer", "a", "good"))

入力

"lina is a good girl she is a good singer"
"a"
"good"

出力

['girl', 'singer']

処理の流れと計算量

このコードでは、まず split(" ") によってテキストを単語ごとに分割します。その後、各インデックス i について、text[i] が first、text[i + 1] が second と一致しているかを確認します。条件を満たしていれば、その直後の単語 text[i + 2] を結果リストに追加します。範囲外参照を防ぐため、i + 2 < len(text) のチェックを必ず行っている点がポイントです。

計算量は、テキストの単語数を n とすると、時間計算量は O(n)、空間計算量も結果リストを含めて O(n) となり、非常に効率的です。