2021-02-08に更新

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

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

(※ 2016-03-03 に書いた記事のクロスポストです)

参考にしたもの

  • Teaching Materials - ここに置いてある「Burrows-Wheeler Transform and FM Index」というタイトルの PDF(ジョンズ・ホプキンス大学 Ben Langmead さんの講義資料)

最初の取っ掛かりとして分りやすかったのがこれ。要点を押さえた簡潔な図と 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((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 まわりを微修正しました。
Originally published at memo88.hatenablog.com
ツイッターでシェア
みんなに共有、忘れないようにメモ

sonota88

Crieitは誰でも投稿できるサービスです。 是非記事の投稿をお願いします。どんな軽い内容でも投稿できます。

また、「こんな記事が読みたいけど見つからない!」という方は是非記事投稿リクエストボードへ!

有料記事を販売できるようになりました!

こじんまりと作業ログやメモ、進捗を書き残しておきたい方はボード機能をご利用ください。
ボードとは?

コメント