Mae向きなブログ

Mae向きな日記のブログ版。ようやくこちらに移行してきました。

検索システム

今日は、「転置インデックスによる検索システムを作ってみよう!」を参考に、学習。

検索システムというと難しい印象がありましたが、上記サイトは簡単なサンプルで解説がされており非常に参考になりました。

サンプルはPerlで書かれていたので、理解を深めるためにRubyで書いてみました。

index.rb

num = 0
idx = Hash.new{|h, key| h[key] = [] }

while line = gets()
  line.chomp!
  next unless %r|\A(\d+) (.+)\z| =~ line
  id, c = $1, $2
  char = []
  c.scan(/\w/){|matched| char << matched }
  seen = Hash.new
  char.each_with_index do |ele, i|
    break if i == char.size-1
    bigram = char[i..i+1].to_s 
    next if seen[bigram]
    idx[bigram] << id
    seen[bigram] = 1
  end
  num += 1
end

puts "\#NUM = #{num}"
idx.sort_by{|key, value| key}.each do |ele|
  puts ele[0] + ' ' + ele[1].join(',')
end

search.rb

num = 0
idx = Hash.new{|h, key| h[key] = [] }

File.open(ARGV[0]) do |f|
  while line = f.gets()
    line.chomp!
    if %r|\A\#NUM = (\d+)\z| =~ line
      num = $1.to_f
    elsif %r|\A(\S+) (.+)\z| =~ line
      idx[$1] = $2.split(/,/)
    end
  end
end

line =  ARGV[1]
score = Hash.new(0)
tf = Hash.new(0)
char = []
line.scan(/\w/){|matched| char << matched }
char.each_with_index do |ele, i|
  break if i == char.size-1
  tf[char[i..i+1].to_s] += 1
end

tf.each do |key, value|
  if idx[key].size != 0
    df = idx[key].size
  else
    df = 0
  end
  idf = Math::log(num / (df + 1))
  tfidf = tf[key] * idf
  idx[key].each do |ele|
    score[ele] += tfidf
  end
end
puts line
score.sort{|a, b|
  b[1] <=> a[1]
}.each{|key, value|
  puts "ID:#{key} SCORE:#{value}"
}

実行結果

mmasa@debian:~/work/ruby/search$ ruby -Ke index.rb test.txt  > test.idx
mmasa@debian:~/work/ruby/search$ ruby -Ke search.rb test.idx '最近ペンギンが好きです。'
最近ペンギンが好きです。
ID:3 SCORE:3.70130197411249
ID:2 SCORE:0.8754687373539
ID:5 SCORE:0.693147180559945
ID:1 SCORE:0.587786664902119
ID:6 SCORE:0.587786664902119
ID:4 SCORE:0.182321556793955