Trie木の純粋な実装とLOUDSのサイズ比較(Java)

本記事は ZOZO Advent Calendar 2021 の2日目の記事です

こんにちは、pakioです。
今回の記事は、過去自分が何度か目を背けて来たTrie木について理解しつつ、効率的な実装であるLOUDSも同時に理解しようとした備忘録です。
以下を参考に実装しています。
簡潔データ構造 LOUDS の解説（全12回、練習問題付き）

今回のソースコードは、全て以下のリポジトリにて公開されています。
https://github.com/pakio/Trie-LOUDS

Trie木について

Trie木はプレフィックスマッチの検索で利用できる順序付きの木構造で、正規表現などの検索と比較して検索の効率化が行えるデータ構造です。
文字列を1文字1頂点から構成された木として考え、複数の木のうち先頭から共通する頂点を共有することで対象を減らし、先頭からの検索を効率的にしています。

詳しい説明は以下の記事などを御覧ください。
https://algo-logic.info/trie-tree/

TrieのJavaでの実装例

JavaでのTrie木の更新系実装例です。

Trie.java

public class Trie {
  private Node rootNode;

  public Trie() {
    rootNode = new Node('-');
  }

  public void add(String s) {
    add(s.toCharArray(), 0, rootNode);
  }

  private static void add(char[] chars, int offset, Node node) {
    char key = chars[offset];

    Node childNode;
    if (node.keyExists(key)) {
      childNode = node.getChildNode(key);
    } else {
      childNode = new Node(key);
      node.add(childNode);
    }

    if (chars.length == offset + 1) {
      childNode.setIsLeaf();
      return;
    }

    add(chars, ++offset, childNode);
  }
}

Node.java

import java.util.HashMap;
import java.util.Map;

public class Node {
  private final char key;
  private boolean isLeaf;
  private Map<Character, Node> childNodes;

  public Node(char k) {
    key = k;
    childNodes = new HashMap<>();
  }

  public char getKey() {
    return key;
  }

  public void setIsLeaf() {
    isLeaf = true;
  }

  public boolean isLeaf() {
    return isLeaf;
  }

  public void add(Node childNode) {
    childNodes.put(childNode.getKey(), childNode);
  }

  public boolean hasChild() {
    return !childNodes.isEmpty();
  }

  public boolean keyExists(char c) {
    return this.childNodes.containsKey(c);
  }

  public Map<Character, Node> getChildNodes() {
    return childNodes;
  }

  public Node getChildNode(char c) {
    return childNodes.get(c);
  }
}

LOUDSについて

LOUDS(Level-Ordered Unary Degree Sequenece)は、木構造をより効率的なメモリ量で実装することを目的とした簡潔データ構造です。
オンラインでの更新が必要ない場合(=木の構造が固定である場合)、木からLOUDSに変換することによって省スペース化が可能です。
純粋に木構造を再現するだけであれば以下の記事で紹介されている通り純粋なビット列での表現だけで表現可能ですが、今回は検索のためそれぞれの頂点が対応する文字、頂点が葉かどうかのフラグも合わせて保持します。
https://takeda25.hatenablog.jp/entry/20120421/1335019644

LOUDSのJavaでの実装例

JavaでTrie木を変換にLOUDS変換するための最低限の実装です。

LOUDS.java

import java.util.ArrayList;
import java.util.List;

public class LOUDS {
  List<Boolean> LBS;
  List<Character> labels;
  List<Boolean> isLeaf;

  public LOUDS () {
    LBS = new ArrayList<>();
    LBS.add(true);
    LBS.add(false);

    labels = new ArrayList<>();
    labels.add(null);
    labels.add(null);

    isLeaf = new ArrayList<>();
    isLeaf.add(false);
    isLeaf.add(false);
  }
}

Converter.java

import java.util.ArrayDeque;
import java.util.Map.Entry;
import java.util.Queue;

public class Converter {
  /**
   * convert tree to LOUDS
   */
  public static LOUDS convert(Node rootNode) {
    LOUDS louds = new LOUDS();
    Queue<Node> queue = new ArrayDeque<>();
    queue.add(rootNode);

    while(!queue.isEmpty()) {
      processQueue(queue, louds);
    }

    return louds;
  }

  public static void processQueue(Queue<Node> queue, LOUDS louds) {
    Node node = queue.poll();

    if (node.hasChild()) {
      for(Entry<Character, Node> characterNodeMap : node.getChildNodes().entrySet()) {
        queue.add(characterNodeMap.getValue());
        louds.LBS.add(true);
        louds.labels.add(characterNodeMap.getKey());
        louds.isLeaf.add(characterNodeMap.getValue().isLeaf());
      }
    }

    // end of node
    louds.LBS.add(false);
    louds.isLeaf.add(false);
  }
}

サイズ比較

LOUDSのメリットはサイズが小さくなることと記載しましたが、本当に小さくなっているのか、素人実装ながら計測してみました。
計測にはjava-sizeofを用います。

計測はApache License2.0の全文をパースした文字列にて行います。

計測結果は以下のようになりました。

実行結果

Trie size(byte) -> 430328
LOUDS size(byte) -> 62384

7分の1程度で表現ができていそうです。

計測に利用したコードは以下のとおりです。

Main.java

import com.carrotsearch.sizeof.RamUsageEstimator;
import com.google.common.base.Charsets;
import com.google.common.io.Resources;
import java.io.IOException;
import java.net.URL;
import java.util.StringTokenizer;

public class Main {
  public static void main(String[] args) throws IOException {
    Trie trie = new Trie();

    URL sourceFile = Resources.getResource("LICENSE.txt");
    String targetString = Resources.toString(sourceFile, Charsets.UTF_8);
    String delim = " \n\r\t,.;";
    StringTokenizer st = new StringTokenizer(targetString,delim);
    while (st.hasMoreTokens()) {
      trie.add(st.nextToken());
    }

    LOUDS louds = trie.convert();

    System.out.println("Trie size(byte) -> " + RamUsageEstimator.sizeOf(trie));
    System.out.println("LOUDS size(byte) -> " + RamUsageEstimator.sizeOf(louds));
  }
}

検索ロジックの実装

ついでに、LOUDSを用いた検索ロジックも実装しました。

LOUDS.java

import java.util.ArrayList;
import java.util.List;

public class LOUDS {
  List<Boolean> LBS;
  List<Character> labels;
  List<Boolean> isLeaf;

  public LOUDS () {
    LBS = new ArrayList<>();
    LBS.add(true);
    LBS.add(false);

    labels = new ArrayList<>();
    labels.add(null);
    labels.add(null);

    isLeaf = new ArrayList<>();
    isLeaf.add(false);
    isLeaf.add(false);
  }

  public boolean match(String s) {
    return search(2, s.toCharArray(), 0);
  }

  private boolean search(int index, char[] chars, int wordOffset) {
    int charIndex = countTrue(index);
    while(LBS.get(index)) {
      if (chars[wordOffset] == labels.get(charIndex)) {
        if (isLeaf.get(index) && wordOffset + 1 == chars.length) return true;
        else if (wordOffset + 1 == chars.length) return false;
        return search(indexOfLabel(charIndex), chars, ++wordOffset);
      } else {
        index ++;
      }
      charIndex ++;
    }
    return false;
  }

  private int countTrue(int to) {
    return (int)LBS.subList(0, to + 1).stream().filter(elm -> elm).count();
  }

  private int indexOfLabel(int label) {
    int count = 0, i = 0;
    for(; i < LBS.size(); i ++) {
      if (!LBS.get(i)) {
        if (++count == label) {
          break;
        }
      }
    }

    return i + 1;
  }
}

検索結果は以下の通りです。

実行結果

match "one" -> true
match "onee" -> false

実行コードは以下の通りです。

Main.java

import com.google.common.base.Charsets;
import com.google.common.io.Resources;
import java.io.IOException;
import java.net.URL;
import java.util.StringTokenizer;

public class Main {
  public static void main(String[] args) throws IOException {
    Trie trie = new Trie();

    URL sourceFile = Resources.getResource("LICENSE.txt");
    String targetString = Resources.toString(sourceFile, Charsets.UTF_8);
    String delim = " \n\r\t,.;";
    StringTokenizer st = new StringTokenizer(targetString,delim);
    while (st.hasMoreTokens()) {
      trie.add(st.nextToken());
    }

    LOUDS louds = trie.convert();

    System.out.println("match \"one\" -> " + louds.match("one"));
    System.out.println("match \"onee\" -> " + louds.match("onee"));
  }
}

まとめ

JavaでTrie木のLOUDSを用いて実装しました。変換処理が都度必要なため若干手間はありそうですが、コンパクトに木が実装できることはかなりメリットが大きそうに感じています。
本来であればここからDFSを用いてサジェスト機能まで実装したかったところですが、今回叶わずだったのでまたの機会に実装してみます。