suffix array構築のメモリ効率を良くする - アルゴリズム学習(その7)

blog.shibayu36.org

上の記事で、一番簡単なアルゴリズムでのsuffix arrayの構築を実装してみた。しかしこれをベンチマークしようとして、10万文字くらいの文字列に対して適応してみると、Java heap spaceというエラーが出てしまい、計算できなかった。

こうなる理由は、上記のアルゴリズムのメモリ効率が非常に悪いからである。メモリ効率が悪い理由は、部分文字列を作るのに文字列コピーが発生しているからだ。

例えば「banana」という文字列なら、前回の記事の実装でsuffix arrayを計算しようとすると、

banana
anana
nana
ana
na
a

という部分文字列をコピーして作ってしまう。このため、1文字を2バイトとすると、21文字 * 2 = 42バイト使ってしまう。つまりメモリ使用量は文字数をnとするとn(n+1)バイト使うことになる。

このメモリ使用量で、10万文字のsuffix arrayを構築しようとすると、100000 * 100001 = 10ギガバイトのメモリを食うことになり、JVMに割り当てられているメモリ量を超えてしまっていたようだ。

解決法

結局は文字列コピーするのが悪いため、コピーはせずに渡された文字を共通に利用していけば良い。ググっていたら正にそういうことをしている例が http://algs4.cs.princeton.edu/63suffix/SuffixArray.java.html にあったので、写経して書いてみる。

まず、全体文字列(text)と、どこからのsuffixなのか(index)を保持するSuffixクラスというものを用意する。

src/main/java/Suffix.java

public class Suffix implements Comparable<Suffix> {
  public String text; // 全体文字列
  public int index;   // 全体文字列中のどこからのsuffixか

  public Suffix(String text, int index) {
    this.text = text;
    this.index = index;
  }
  public int length() {
    return text.length() - index;
  }
  public char charAt(int i) {
    return text.charAt(index + i);
  }

  public int compareTo(Suffix that) {
    if (this == that) return 0;  // optimization
    int n = Math.min(this.length(), that.length());
    for (int i = 0; i < n; i++) {
      if (this.charAt(i) < that.charAt(i)) return -1;
      if (this.charAt(i) > that.charAt(i)) return +1;
    }
    return this.length() - that.length();
  }
}

このクラスはComparableを実装しているので、Suffixクラス同士で比べることができるようになっている。compareToが比べるためのメソッドで、先頭から1文字ずつ比較しながら辞書順に並べられるようにしている。

そしてこのクラスを用いて、SuffixArrayの構築を行う。

src/main/java/SuffixArray2.java

import java.util.Arrays;

public class SuffixArray2 {
  public static Integer[] make(String text) {
    int n = text.length();
    Suffix[] suffixes = new Suffix[n];
    for (int i = 0; i < n; i++) {
      suffixes[i] = new Suffix(text, i);
    }
    Arrays.sort(suffixes);
    return Arrays.stream(suffixes).map(s -> s.index).toArray(size -> new Integer[size]);
  }
}

これで完成。textを使いまわしていて、String.substringを使っていないため、文字列のコピーが発生せず、メモリをtext * 2バイト以上使うことはない。

まとめ

今回は前回のsuffix array実装を、メモリ効率の良い実装に変えてみた。今回はアルゴリズム自体を変えているわけではないので、計算量は変わらず ${\mathcal{O} ( n^2 \log n ) }$ ではあるが、文字長が長くなったとしてもJava heap sizeなどのエラーを出さずに実行ができるようになった。