Scala で転置インデックスによる検索システム

「転置インデックスによる検索システムを作ってみよう！」 [1] という記事があったのでそれを真似して Scala で書いてみました。

import java.io._
import scala.io._
import scala.collection.immutable._

object Bigram {
  type Bigram = (Char, Char)
  def bigrams(s: String) = s.toList.zip(s.toList.drop(1))
}
import Bigram._

object Index {
  def save(file: String, obj: AnyRef) = new ObjectOutputStream(new FileOutputStream(file)).writeObject(obj)
  def main(args: Array[String]) {
    val Array(srcfile, idxfile) = args
    val lines = Source.fromFile(srcfile).getLines
    var idx = new TreeMap[Bigram, Set[Int]]
    var n = 0
    for (line <- lines) {
      n += 1
      for (bi <- bigrams(line.stripLineEnd)) {
        val newval = idx.get(bi) match {
          case Some(s) => s + n
          case None    => TreeSet(n)
        } // withDefaultValue を使うとシリアライズできなくなるので
        idx = idx.update(bi, newval)
      }
    }
    save(idxfile, (n, idx))
  }
}

object Search {
  def load(file: String) = new ObjectInputStream(new FileInputStream(file)).readObject
  def main(args: Array[String]) {
    val Array(key, idxfile) = args
    val (num, idx) = load(idxfile).asInstanceOf[(Int, TreeMap[Bigram, Set[Int]])]
    val idx_d = idx.withDefaultValue(new TreeSet[Int])
    val bi_key = bigrams(key)
    // tf: Bigram => Int
    val tf = bi_key.foldLeft(new TreeMap[Bigram, Int].withDefaultValue(0)) {
      (tf, bi) => tf.update(bi, tf(bi) + 1)
    }
    def df(bi: Bigram): Int = idx_d(bi).size
    def w(bi: Bigram): Double = tf(bi) * Math.log(num.toDouble / (df(bi) + 1))

    val scores =
      for (linenum <- (1 to num).toList)
        yield (bi_key.foldLeft(0.0) {
          (sum, bi) => if (idx_d(bi).contains(linenum)) sum + w(bi) else sum
        }, linenum)

    for ((score, linenum) <- scores.sort{(l,r) => (l compare r) > 0}) {
      println("ID:" + linenum + " SCORE:" + score)
    }
  }
}

元記事から以下の点を変えました。

1. 元記事ではソースファイルに記事IDという列があったのを省いて単に行番号を付与した
2. インデックスファイルの作成には Java のシリアライザを使った
3. 検索プログラムのほうで [2] の記事に出てくる数式に対応する形ができるだけコード中にも出る感じにした
4. 全体に標準入出力は使わずコマンドラインオプションで指定にした

ファイル閉じていないとかは大目に見てください。

実行結果：

PS D:\scala> cat test.txt
これはペンです
最近はどうですか？
ペンギン大好き
こんにちは。いかがおすごしですか？
ここ最近疲れ気味
ペンキ塗りたてで気味が悪いです
PS D:\scala> scala Index test.txt test.idx
PS D:\scala> scala Search "最近ペンギンが好きです" test.idx
ID:3 SCORE:3.7013019741124933
ID:2 SCORE:0.8754687373538999
ID:5 SCORE:0.6931471805599453
ID:6 SCORE:0.587786664902119
ID:1 SCORE:0.587786664902119
ID:4 SCORE:0.1823215567939546

[1] http://chalow.net/2007-11-26-5.html
[2] http://chalow.net/2005-10-12-1.html