雲梯をすすむように

Ruby/RBTreeの存在をいままでしりませんでしたが、これ、おもろい。RB木（SBB木?）を提供する拡張ライブラリで、当然だけどRubyで書いた誰かのSBB木よりも相当速いと予想されます。

例によって転置インデックスをつくるよ。

  require 'rbtree'

  tree = RBTree.new
  while line = ARGF.gets
    line.scan(/(\w+)/) do |words|
      pos = Regexp.last_match.begin(0)
      key = [words[0], ARGF.path, ARGF.file.lineno, pos]
      tree[key] = true
    end
  end

単語、ファイル名、行番号、行内の位置のArrayをkeyとします。valueはどうでもいいのでtrueで。

RBTreeのboundメソッドを使うとある範囲の要素を順に辿れます。'def'が含まれる場所を調べるにはこんな感じ。

  tree.bound(['def'], tree.last[0]).each do |k, v|
    break unless k[0] == 'def'
    p k
  end

実行するとこんなの。

["def", "rinda_eval.rb", 6, 4]
["def", "rinda_eval.rb", 12, 4]
["def", "rinda_eval.rb", 18, 2]
["def", "tree.rb", 5, 4]
["def", "tree.rb", 12, 4]
["def", "tree.rb", 19, 2]
["def", "tree.rb", 47, 15]
["def", "tree.rb", 48, 26]

ほかにlower_boundというメソッドがあって、指定したキーの要素（なければその次のキーの要素）を調べることができます。それをつかってand検索の練習をします。

ある単語群が同時に現れる行を返す例を示します。
# Searchのeachが引数をとるのがイマイチだけど、よい名前が思いつかなかったので我慢。
Search::Cursorは単語ごとに生成されるカーソルで、lower_boundを使って実装されます。指定されたファイル名、行番号、行内位置、あるいはその次の場所を調べる係です。

require 'rbtree'

class Search
  class Cursor
    def initialize(tree, word)
      @tree = tree
      @word = word
      @cur = nil
    end
    attr_reader :cur
    
    def forward(path='', line=0, pos=0)
      key, value = @tree.lower_bound([@word, path, line, pos])
      throw(:cursor_eof) unless key && key[0] == @word
      @cur = key[1, 3]
    end
  end

  def each(tree, words)
    cursor = words.collect {|w| Cursor.new(tree, w)}
    catch(:cursor_eof) do
      cursor.each {|c| c.forward}
      while true
        sorted = cursor.sort_by {|c| c.cur}
        path, lineno, pos = sorted.last.cur
        if sorted.first.cur[0, 2] == [path, lineno]
          yield(path, lineno)
          cursor.each {|c| c.forward(path, lineno + 1)} # next lineno
        else
          sorted.first.forward(path, lineno)
        end
      end
    end
  end
end

if __FILE__ == $0
  tree = RBTree.new
  while line = ARGF.gets
    line.scan(/(\w+)/) do |words|
      pos = Regexp.last_match.begin(0)
      key = [words[0], ARGF.path, ARGF.file.lineno, pos]
      tree[key] = true
    end
  end

  words = %w(unless return)
  Search.new.each(tree, words) do |path, lineno|
    p [path, lineno]
  end

上の例では'unless'と'return'が同時に出現する行を調べてます。

eachがミソなのでも一度載せます。

• 単語ごとにカーソルをつくる
• 現在位置でソート
• 先端の行と末尾の行が一致してたらマッチ。全てのカーソルを一行すすめる。
• マッチしてなかったら、末尾のカーソルを先端のカーソルと同じ行へすすめる。
• だれかが尽きたら終わり

雲梯にぶらさがりながら、腕を振ってすすむようにすすむです。

  def each(tree, words)
    cursor = words.collect {|w| Cursor.new(tree, w)}
    catch(:cursor_eof) do
      cursor.each {|c| c.forward}
      while true
        sorted = cursor.sort_by {|c| c.cur}
        path, lineno, pos = sorted.last.cur
        if sorted.first.cur[0, 2] == [path, lineno]
          yield(path, lineno)
          cursor.each {|c| c.forward(path, lineno + 1)} # next lineno
        else
          sorted.first.forward(path, lineno)
        end
      end
    end
  end

なお、この辺りを参考にしましたです。

• 転置インデックスの構成とブーリアン検索