ma.vi.base.trie.Trie Maven / Gradle / Ivy

Go to download
package ma.vi.base.trie;

import ma.vi.base.tuple.T2;

import java.util.*;

import static java.util.Collections.emptyList;

/**
 * A trie keeps sequences of a unit type U mapped to a value
 * type V. A trie of unit type Character, for instance, would
 * store sequences of characters (i.e. strings) mapped to
 * objects of some value type.
 *
 * @author Vikash Madhow ([email protected])
 */
public class Trie {

  public V put(Iterable sequence, V value) {
    Node node = root;
    for (U u: sequence) {
      if (node.children.containsKey(u)) {
        node = node.children.get(u);
      } else {
        Node child = new Node<>();
        node.children.put(u, child);
        node = child;
      }
    }
    V previous = node.value;
    node.value = value;
    return previous;
  }

  public V get(Iterable sequence) {
    Node node = find(sequence.iterator());
    return node == null ? null : node.value;
  }

  public List, V>> getPrefixed(Iterable sequence) {
    Node node = root;
    List prefix = new ArrayList<>();
    for (U u: sequence) {
      prefix.add(u);
      if (node.children.containsKey(u)) {
        node = node.children.get(u);
      } else {
        return emptyList();
      }
    }
    return getPrefixed(prefix, node);
  }

  public void delete(Iterable sequence) {
    delete(sequence.iterator(), root);
  }

  private boolean delete(Iterator sequence, Node node) {
    if (sequence.hasNext()) {
      U u = sequence.next();
      if (node.children.containsKey(u)) {
        boolean deletedAll = delete(sequence, node.children.get(u));
        if (deletedAll) {
          node.children.remove(u);
        }
      }
    } else {
      node.value = null;
    }
    return node.children.isEmpty();
  }

  public boolean deletePrefixed(Iterable sequence) {
    Node node = find(sequence.iterator());
    if (node != null) {
      node.children.clear();
    }
    return false;
  }

  protected Node find(Iterator sequence) {
    Node node = root;
    while (sequence.hasNext()) {
      U u = sequence.next();
      if (node.children.containsKey(u)) {
        node = node.children.get(u);
      } else {
        return null;
      }
    }
    return node;
  }

  private List, V>> getPrefixed(List prefix, Node startFrom) {
    List, V>> values = new ArrayList<>();
    if (startFrom.children.isEmpty()) {
      List list = new ArrayList<>(prefix);
      values.add(T2.of(list, startFrom.value));
    } else {
      for (Map.Entry> child: startFrom.children.entrySet()) {
        List newPrefix = new ArrayList<>(prefix);
        newPrefix.add(child.getKey());
        values.addAll(getPrefixed(newPrefix, child.getValue()));
      }
    }
    return values;
  }

  private static class Node {
    final Map> children = new HashMap<>();
    V value;
  }

  private final Node root = new Node<>();
}