【python】zipを使ってn-gram列を生成する

はじめに

　n-gramは自然言語処理でよく使われる方法です。

n-gram - Wikipedia

　さて、以下のような関数を作りたいとします。

n_gram("abcde", n=2, sep="-")  # ["a-b", "b-c", "c-d", "d-e"]

　n=2ならbigram, n=3ならtrigramという言い方があります。さて、たとえばbigramなら以下のように書けます。

>>> def bigram(seq, sep="-"):
...     return [sep.join(x) for x in zip(seq, seq[1:])]
... 
>>> bigram("abcde")
['a-b', 'b-c', 'c-d', 'd-e']

　これは私の発明ではありませんが、pythonに慣れている人なら誰でも思いつく可能性はあると思います。

　trigramはこう。

>>> def trigram(seq, sep="-"):
...     return [sep.join(x) for x in zip(seq, seq[1:], seq[2:])]
... 
>>> trigram("abcde")
['a-b-c', 'b-c-d', 'c-d-e']

一般化

　zipの引数部分をnの数だけ繰り返せば、一般化することができます。ここはジェネレータ式の出番です。

>>> def n_gram(seq, n=2, sep="-"):
...     return [sep.join(x) for x in zip(*(seq[i:] for i in range(n)))]
... 
>>> n_gram("abcde", n=1)
['a', 'b', 'c', 'd', 'e']
>>> n_gram("abcde", n=2)
['a-b', 'b-c', 'c-d', 'd-e']
>>> n_gram("abcde", n=3)
['a-b-c', 'b-c-d', 'c-d-e']
>>> n_gram("abcde", n=4)
['a-b-c-d', 'b-c-d-e']
>>> n_gram("abcde", n=5)
['a-b-c-d-e']
>>> n_gram("abcde", n=6)
[]

効率化

　上の実装はラフに見てseqのn倍のメモリ領域を食ってしまいます。

　そんなに深刻な問題でもありませんが、とはいえこのままではなんとなくエレガントではないので、isliceを使いたいと思います。

itertools --- 効率的なループ実行のためのイテレータ生成関数 — Python 3.8.1 ドキュメント

　ドキュメントには明記されていなかったと思いますが、seq[i:]とするにはislice(seq, i, None)でいいのだと思います。

>>> from itertools import islice
>>> def n_gram(seq, n=2, sep="-"):
...     return [sep.join(x) for x in zip(*(islice(seq, i, None) for i in range(n)))]
... 
>>> n_gram("abcde", n=1)
['a', 'b', 'c', 'd', 'e']
>>> n_gram("abcde", n=2)
['a-b', 'b-c', 'c-d', 'd-e']
>>> n_gram("abcde", n=3)
['a-b-c', 'b-c-d', 'c-d-e']
>>> n_gram("abcde", n=4)
['a-b-c-d', 'b-c-d-e']
>>> n_gram("abcde", n=5)
['a-b-c-d-e']
>>> n_gram("abcde", n=6)
[]

シーケンスを第一引数に取る意味

　任意のイテラブルを取れるようにしておくと、

>>> n_gram(["I",  "Am", "a", "Cat"])
['I-Am', 'Am-a', 'a-Cat']

　のように応用が効いたりして便利です。