みずぴー日記

SuffixArray

Ocaml 30分プログラム

30分プログラム、その580。id:Gemmaさんに借りたWEB+DB PRESS Vol.50に、suffix arrayの解説が載っていたのでやってみた。
解説を読んだときは「ちょう簡単じゃん。さくっと実装してやんよ」と思っていたけど、いざ始めたけど、けっこう大変だった。簡単とか言って、ごめんなさい。

そもそもararyとついてる時点で大変なことに気がつくべきだった。ボク、OCamlでarrayを使ったことなんてほとんどないじゃないか。

使い方

シグネチャはこんな感じ。

type t
val make : string -> t
val find : t -> string -> int list

まず、suffix arrayを作る。

# let s =
  SuffixArray.make "abracadabra";;
  val s : SuffixArray.t = <abstr>

で、文字列を検索すると、何文字目から始まっているかが分かる。

# let t =
  SuffixArray.find s "ra";;
  val t : int list = [2; 9]

ソースコード

let (@@) f g = f g

module SuffixArray : sig
  type t
  val make : string -> t
  val find : t -> string -> int list
end = struct
  type t = (string*int) array
  let make s =
    let len =
      String.length s in
    let array =
      Array.init len (fun i -> (String.sub s i (len - i),i)) in
      Array.stable_sort (fun (a,_) (b,_) -> compare a b) array;
      array

  let begin_with (entry,_) subject =
    let subject_size =
      String.length subject in
    let entry_size =
      String.length entry in
      if entry_size < subject_size then
	false
      else
	(String.sub entry 0 subject_size) = subject

  let rec take_back f array i =
    if i < Array.length array && f array.(i) then
      array.(i)::take_back f array (i+1)
    else
      []

  let rec take_forward f array i =
    if 0 <= i && f array.(i) then
      array.(i)::take_forward f array (i-1)
    else
      []

  let take_near f array i =
    array.(i)::take_back f array (i+1) @ take_forward f array (i-1)

  let rec find' s subject first last =
    if first >= last then
      []
    else
      let mid =
	(first + last) / 2 in
	if begin_with s.(mid) subject then
	  List.map snd @@ take_near (fun x -> begin_with x subject) s mid
	else
	    if (fst s.(mid)) < subject then
	      find' s subject (mid+1) last
	    else
	      find' s subject first mid

  let find s subject =
    find' s subject 0 (Array.length s)
end

let s =
  SuffixArray.make "abracadabra";;
let t =
  SuffixArray.find s "ra"

参考