Pythonで転置インデックスによる検索システム
山下たつをさんの転置インデックスによる検索システムを作ってみよう!のコードをPythonで真似してみました。Python2.5じゃないと動きません。
ファイルフォーマットや使い方はそのままですが、フォーマット検査をちゃんとやってないです。正規表現のキャプチャ、Pythonでも簡単に書けないかなぁ。
index.py
#!/usr/bin/env python import sys, codecs from collections import defaultdict sys.stdout = codecs.getwriter("utf-8")(sys.stdout) index = defaultdict(lambda:list()) num_docs = 0 for line in sys.stdin: doc = line.decode("utf-8").strip().split(" ", 1) if len(doc) != 2: continue id, text = doc bigram_set = set() for i in xrange(len(text)-1): bigram = "".join((text[i], text[i+1])) if bigram in bigram_set: continue index[bigram].append(id) bigram_set.add(bigram) num_docs += 1 print "#NUM=%s" % num_docs for key in sorted(index.keys()): print "%s %s" % (key, ",".join(index[key]))
search.py
#!/usr/bin/env python from __future__ import with_statement import sys, math from collections import defaultdict index = dict() num_docs = 0 with file(sys.argv[1]) as f: for line in f: line = line.decode("utf-8").strip() if line.startswith("#NUM="): num_docs = float(line.split("=")[-1]) else: bigram, ids = line.split(" ") index[bigram] = ids.split(",") for line in sys.stdin: line = line.decode("utf-8").strip() score_dict = defaultdict(float) tf_dict = defaultdict(int) for i in xrange(len(line)-1): tf_dict["".join((line[i], line[i+1]))] += 1 for bigram, tf in tf_dict.iteritems(): df = len(index[bigram]) if (bigram in index) else 1 idf = math.log(num_docs/(df+1)) tf_idf = tf * idf for id in index.get(bigram, ()): score_dict[id] += tf_idf for id, score in sorted(score_dict.items(), lambda x,y: cmp(y[1], x[1])): print "ID:%s SCORE:%s" % (id, score)
実行
$ cat test.txt
1 これはペンです
2 最近はどうですか?
3 ペンギン大好き
4 こんにちは。いかがおすごしですか?
5 ここ最近疲れ気味
6 ペンキ塗りたてで気味が悪いです
$ ./index.py test.txt > test.idx
$ echo '最近ペンギンが好きです' | ./search.py test.idx
ID:3 SCORE:3.70130197411
ID:2 SCORE:0.875468737354
ID:5 SCORE:0.69314718056
ID:1 SCORE:0.587786664902
ID:6 SCORE:0.587786664902
ID:4 SCORE:0.182321556794
最初、logの前の除算をintどうしでやってしまいスコアの精度がやたら低くなってしまいました。