小ワザ : zip(*[lists])を用いてN-gramを作成する

この記事は・・・

zip(*[lists])の挙動の理解に時間がかかった為、備忘録として書きました。

はじめに

現在、自然言語処理系のKaggleコンペに取り組んでいます。
コンペ参加者のノートブックの中に、探索的データ分析(Exploratory Data Analysis : EDA)を行う目的で N-gramのリストを作成する関数を発見しました。

from wordcloud import STOPWORDS # STOPWORDSのリストを取得

def generate_ngrams(text, n_gram=1):
    token = [token for token in text.lower().split(' ') if token != '' if token not in STOPWORDS]
    ngrams = zip(*[token[i:] for i in range(n_gram)])
    return [' '.join(ngram) for ngram in ngrams]

# 試しに generate_ngrams()を使ってみる
text = "I wanna be a deep python engineer"
ngrams = generate_ngrams(text, n_gram=2)
print(ngrams)

実行結果

['wanna deep', 'deep python', 'python engineer']

なるほど、引数n_gramを変えることで、引数textをトークナイズし、ストップワードは取り除いた上でリスト形式のN-gramのリストを返してくれる関数みたいですね。
なんとなくそのままスルーしそうになりましたが、ngrams = zip(*[token[i:] for ...の一行が気になりました。
少し考えてみましたが、直感的に理解できなかった。。。のでノートブックで色々トライして理解してみることにしました。

結論

* 演算子をつけることで、各リストから、先頭要素から順番に、同じインデックスの要素を取り出して新しいリストを作ることができます。
公式ドキュメントのzip(*iterables)の項目によると、

zip() に続けて * 演算子を使うと、zip したリストを元に戻せます:

つまり、unzipを行っていることになるそうです。
また、同様のテーマを取り上げたQiitaの記事にもあるように、「行と列を変換している」とイメージすると直感的に理解しやすいかもしれません。

「N-gramを作成する」という目的で、zip(*[lists])を使うとするならば、中身のリストは、トークンのリストから先頭の要素を1つずつ削除したリストをn個並べるリストだと良さそうです。
上記の関数では[token[i:] for i in range(n_gram)]が正に目的のリストを作成しています。
綺麗なコードだなと感心しました。

理解のためにトライしたこと

理解のために実行したコードをだらだらと記しています。。。

zip(*[lists])のlists部分を正しく理解してみた。

token = ["wanna", "deep", "python", "engineer"] # 関数実行で得られる tokenのリスト
n_gram = 3 # トライグラムのコーパスを作る
lists = [token[i:] for i in range(n_gram)]
print(lists)

実行結果

[['wanna', 'deep', 'python', 'engineer'], ['deep', 'python', 'engineer'], ['python', 'engineer']]

分かりづらいのでzip(*[lists])の中身のリスト部分を重複しない数字に置き換えてみた。

for nums in zip(*[1,2,3,4], [5,6,7], [8,9]):
    print(nums)

実行結果

(1, 5, 8)
(2, 6, 9)

3つのリストの1番目要素を取り出したリスト、2番目要素を取り出したリストが作られました。