Rubyで素朴なFM-indexを書いてみた

Ruby

BWT、検索処理の最適化・高速化は行なっていません（SA-IS、ウェーブレット行列などは使っていません）。
BWT から検索まで全体の流れが見渡せる最小限の実装にしました。
せっかくなので Ruby に馴染みのない方が見ても読みやすいと思われる書き方にしています（returnやメソッド呼び出しの括弧を省略しない、など）。

参考にしたもの

最初の取っ掛かりとして分りやすかったのがこれ。要点を押さえた簡潔な図と Python コードを眺めてるだけでだいぶ分かった気になれます。


もっと具体的なところについては実装を見た方が早いということで ハクビシンにもわかる全文検索 の erukiti さんの実装を参考にさせてもらいました。CoffeeScript 製。

コード

- range
  - (a..b)  => [a, b] b is included
  - (a...b) => [a, b) b is not included
  - a.downto(b) => [a, a-1, ... b+1, b]
- method{|x| ... } => in JavaScript: method(function(x){ ... })
require "minitest/autorun"

SENTINEL = "$"

def to_bwm(t)
  tt = t + SENTINEL + t
  rows = (0..t.length).map {|i|
    tt[i..(i + t.length)]
  }
  return rows.sort
end

# Burrows-Wheeler transform
def bwt(t)
  bwm = to_bwm(t)
  return bwm.map {|cs| cs[-1] }.join("")
end

def rank_less_than(cs, c)
  return cs.count {|_c| _c < c }
end

def rank(cs, c, i)
  return (0...i).count {|j| cs[j] == c }
end

# LF mapping
def map_lf(cs, c, i)
  return rank_less_than(cs, c) + rank(cs, c, i)
end

# String before c
def backward_chars(bwt_t, c, s, e)
  bwt_cs = bwt_t.split("") # to array of chars

  cb = c   # T(i)
  ca = nil # T(i-1)
  result = ""

  while true
    s = map_lf(bwt_cs, cb, s)
    e = map_lf(bwt_cs, cb, e)
    # assert e - s == 1
    ca = bwt_t[s]
    if (ca == SENTINEL)
      break
    end
    result += ca
    cb = ca
  end

  return result.reverse
end

def reverse(bwt_t)
  # sentinel は F列では必ず 0 行目のみに存在する
  # sentinel is always at F[0]
  s = 0 # start
  e = 1 # end

  return backward_chars(bwt_t, SENTINEL, s, e)
end

def search_internal(bwt_t, q)
  bwt_cs = bwt_t.split("") # to array of chars
  s = 0
  e = bwt_t.length

  (q.length - 1).downto(0).each {|i|
    s = map_lf(bwt_cs, q[i], s)
    e = map_lf(bwt_cs, q[i], e)
    if (s >= e)
      return nil
    end
  }

  return [s, e]
end

def search(bwt_t, q)
  s, e = search_internal(bwt_t, q)
  return s == nil ? 0 : (e - s)
end


class FmIndexTest < Minitest::Test

  def setup
    t = "abaaba"
    @bwt_t = bwt(t)
  end

  def test_bwt
    assert_equal("abba$aa", @bwt_t)
  end

  # 1文字の検索
  def test_one_char
    s, e = search_internal(@bwt_t, "a")
    assert_equal(1, s)
    assert_equal(5, e)

    # 出現回数
    num_hits = search(@bwt_t, "a")
    assert_equal(4, num_hits)
  end

  # 2回出現する
  def test_2_hits
    s, e = search_internal(@bwt_t, "aba")
    assert_equal(3, s)
    assert_equal(5, e)

    assert_equal(2, search(@bwt_t, "aba"))
  end

  # 1回出現する
  def test_one_hit
    s, e = search_internal(@bwt_t, "aaba")
    assert_equal(2, s)
    assert_equal(3, e)

    assert_equal(1, search(@bwt_t, "aaba"))
  end

  # 存在しない組み合わせ
  def non_existent_combination
    s, e = search_internal(@bwt_t, "baba")
    assert_equal(nil, s)
    assert_equal(nil, e)

    assert_equal(0, search(@bwt_t, "baba"))
  end

  # 存在しない文字
  def non_existent_char
    s, e = search_internal(@bwt_t, "x")
    assert_equal(nil, s)
    assert_equal(nil, e)

    assert_equal(0, search(@bwt_t, "x"))
  end

  def test_reverse
    bwt_t = bwt("mississippi")
    assert_equal("ipssm$pissii", bwt_t)
    assert_equal("mississippi", reverse(bwt_t))
  end

end
  • （追記 2021-02-08） Ruby 3.0.0 向けに minitest まわりを微修正しました。

関連