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

Go to download
Show more of this group Show more artifacts with this name
Show all versions of com.vikmad.base Show documentation
Base algos, data structures and utilities
The newest version!
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<>();
}