source: trunk/EventBenchCore/src/de/ugoe/cs/eventbench/models/Trie.java @ 251

package de.ugoe.cs.eventbench.models;

import java.io.Serializable;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import de.ugoe.cs.util.StringTools;

import edu.uci.ics.jung.graph.DelegateTree;
import edu.uci.ics.jung.graph.Graph;
import edu.uci.ics.jung.graph.Tree;

/**
 * <p>
 * This class implements a <it>trie</it>, i.e., a tree of sequences that the
 * occurence of subsequences up to a predefined length. This length is the trie
 * order.
 * </p>
 * 
 * @author Steffen Herbold
 * 
 * @param <T>
 *            Type of the symbols that are stored in the trie.
 * 
 * @see TrieNode
 */
public class Trie<T> implements IDotCompatible, Serializable {

        /**
         * <p>
         * Id for object serialization.
         * </p>
         */
        private static final long serialVersionUID = 1L;

        /**
         * <p>
         * Collection of all symbols occuring in the trie.
         * </p>
         */
        private Collection<T> knownSymbols;

        /**
         * <p>
         * Reference to the root of the trie.
         * </p>
         */
        private final TrieNode<T> rootNode;

        /**
         * <p>
         * Contructor. Creates a new Trie.
         * </p>
         */
        public Trie() {
                rootNode = new TrieNode<T>();
                knownSymbols = new LinkedHashSet<T>();
        }

        /**
         * <p>
         * Returns a collection of all symbols occuring in the trie.
         * </p>
         * 
         * @return symbols occuring in the trie
         */
        public Collection<T> getKnownSymbols() {
                return new LinkedHashSet<T>(knownSymbols);
        }

        /**
         * <p>
         * Trains the current trie using the given sequence and adds all subsequence
         * of length {@code maxOrder}.
         * </p>
         * 
         * @param sequence
         *            sequence whose subsequences are added to the trie
         * @param maxOrder
         *            maximum length of the subsequences added to the trie
         */
        public void train(List<T> sequence, int maxOrder) {
                if( maxOrder<1 ) {
                        return;
                }
                IncompleteMemory<T> latestActions = new IncompleteMemory<T>(maxOrder);
                int i = 0;
                for (T currentEvent : sequence) {
                        latestActions.add(currentEvent);
                        knownSymbols.add(currentEvent);
                        i++;
                        if (i >= maxOrder) {
                                add(latestActions.getLast(maxOrder));
                        }
                }
                int sequenceLength = sequence.size();
                for (int j = maxOrder - 1; j > 0; j--) {
                        add(sequence.subList(sequenceLength - j, sequenceLength));
                }
        }

        /**
         * <p>
         * Adds a given subsequence to the trie and increases the counters
         * accordingly.
         * </p>
         * 
         * @param subsequence
         *            subsequence whose counters are increased
         * @see TrieNode#add(List)
         */
        protected void add(List<T> subsequence) {
                if (subsequence != null && !subsequence.isEmpty()) {
                        knownSymbols.addAll(subsequence);
                        subsequence = new LinkedList<T>(subsequence); // defensive copy!
                        T firstSymbol = subsequence.get(0);
                        TrieNode<T> node = getChildCreate(firstSymbol);
                        node.add(subsequence);
                }
        }

        /**
         * <p>
         * Returns the child of the root node associated with the given symbol or
         * creates it if it does not exist yet.
         * </p>
         * 
         * @param symbol
         *            symbol whose node is required
         * @return node associated with the symbol
         * @see TrieNode#getChildCreate(Object)
         */
        protected TrieNode<T> getChildCreate(T symbol) {
                return rootNode.getChildCreate(symbol);
        }

        /**
         * <p>
         * Returns the child of the root node associated with the given symbol or
         * null if it does not exist.
         * </p>
         * 
         * @param symbol
         *            symbol whose node is required
         * @return node associated with the symbol; null if no such node exists
         * @see TrieNode#getChild(Object)
         */
        protected TrieNode<T> getChild(T symbol) {
                return rootNode.getChild(symbol);
        }

        /**
         * <p>
         * Returns the number of occurences of the given sequence.
         * </p>
         * 
         * @param sequence
         *            sequence whose number of occurences is required
         * @return number of occurences of the sequence
         */
        public int getCount(List<T> sequence) {
                int count = 0;
                TrieNode<T> node = find(sequence);
                if (node != null) {
                        count = node.getCount();
                }
                return count;
        }

        /**
         * <p>
         * Returns the number of occurences of the given prefix and a symbol that
         * follows it.<br>
         * Convenience function to simplify usage of {@link #getCount(List)}.
         * </p>
         * 
         * @param sequence
         *            prefix of the sequence
         * @param follower
         *            suffix of the sequence
         * @return number of occurences of the sequence
         * @see #getCount(List)
         */
        public int getCount(List<T> sequence, T follower) {
                List<T> tmpSequence = new LinkedList<T>(sequence);
                tmpSequence.add(follower);
                return getCount(tmpSequence);

        }

        /**
         * <p>
         * Searches the trie for a given sequence and returns the node associated
         * with the sequence or null if no such node is found.
         * </p>
         * 
         * @param sequence
         *            sequence that is searched for
         * @return node associated with the sequence
         * @see TrieNode#find(List)
         */
        public TrieNode<T> find(List<T> sequence) {
                if (sequence == null || sequence.isEmpty()) {
                        return rootNode;
                }
                List<T> sequenceCopy = new LinkedList<T>(sequence);
                TrieNode<T> result = null;
                TrieNode<T> node = getChild(sequenceCopy.get(0));
                if (node != null) {
                        sequenceCopy.remove(0);
                        result = node.find(sequenceCopy);
                }
                return result;
        }

        /**
         * <p>
         * Returns a collection of all symbols that follow a given sequence in the
         * trie. In case the sequence is not found or no symbols follow the sequence
         * the result will be empty.
         * </p>
         * 
         * @param sequence
         *            sequence whose followers are returned
         * @return symbols following the given sequence
         * @see TrieNode#getFollowingSymbols()
         */
        public Collection<T> getFollowingSymbols(List<T> sequence) {
                Collection<T> result = new LinkedList<T>();
                TrieNode<T> node = find(sequence);
                if (node != null) {
                        result = node.getFollowingSymbols();
                }
                return result;
        }

        /**
         * <p>
         * Returns the longest suffix of the given context that is contained in the
         * tree and whose children are leaves.
         * </p>
         * 
         * @param context
         *            context whose suffix is searched for
         * @return longest suffix of the context
         */
        public List<T> getContextSuffix(List<T> context) {
                List<T> contextSuffix = new LinkedList<T>(context); // defensive copy
                boolean suffixFound = false;

                while (!suffixFound) {
                        if (contextSuffix.isEmpty()) {
                                suffixFound = true; // suffix is the empty word
                        } else {
                                TrieNode<T> node = find(contextSuffix);
                                if (node != null) {
                                        if (!node.getFollowingSymbols().isEmpty()) {
                                                suffixFound = true;
                                        }
                                }
                                if (!suffixFound) {
                                        contextSuffix.remove(0);
                                }
                        }
                }

                return contextSuffix;
        }

        /**
         * <p>
         * Helper class for graph visualization of a trie.
         * </p>
         * 
         * @author Steffen Herbold
         * @version 1.0
         */
        static public class Edge {
        }

        /**
         * <p>
         * Helper class for graph visualization of a trie.
         * </p>
         * 
         * @author Steffen Herbold
         * @version 1.0
         */
        static public class TrieVertex {

                /**
                 * <p>
                 * Id of the vertex.
                 * </p>
                 */
                private String id;

                /**
                 * <p>
                 * Contructor. Creates a new TrieVertex.
                 * </p>
                 * 
                 * @param id
                 *            id of the vertex
                 */
                protected TrieVertex(String id) {
                        this.id = id;
                }

                /**
                 * <p>
                 * Returns the id of the vertex.
                 * </p>
                 * 
                 * @see java.lang.Object#toString()
                 */
                @Override
                public String toString() {
                        return id;
                }
        }

        /**
         * <p>
         * Returns a {@link Graph} representation of the trie.
         * </p>
         * 
         * @return {@link Graph} representation of the trie
         */
        protected Tree<TrieVertex, Edge> getGraph() {
                DelegateTree<TrieVertex, Edge> graph = new DelegateTree<TrieVertex, Edge>();
                rootNode.getGraph(null, graph);
                return graph;
        }

        /*
         * (non-Javadoc)
         * 
         * @see de.ugoe.cs.eventbench.models.IDotCompatible#getDotRepresentation()
         */
        public String getDotRepresentation() {
                StringBuilder stringBuilder = new StringBuilder();
                stringBuilder.append("digraph model {" + StringTools.ENDLINE);
                rootNode.appendDotRepresentation(stringBuilder);
                stringBuilder.append('}' + StringTools.ENDLINE);
                return stringBuilder.toString();
        }

        /**
         * <p>
         * Returns the string representation of the root node.
         * </p>
         * 
         * @see TrieNode#toString()
         * @see java.lang.Object#toString()
         */
        @Override
        public String toString() {
                return rootNode.toString();
        }

        /**
         * <p>
         * Returns the number of symbols contained in the trie.
         * </p>
         * 
         * @return number of symbols contained in the trie
         */
        public int getNumSymbols() {
                return knownSymbols.size();
        }

        /**
         * <p>
         * Returns the number of trie nodes that are ancestors of a leaf. This is
         * the equivalent to the number of states a first-order markov model would
         * have.
         * <p>
         * 
         * @return number of trie nodes that are ancestors of leafs.
         */
        public int getNumLeafAncestors() {
                Set<TrieNode<T>> ancestors = new HashSet<TrieNode<T>>();
                rootNode.getLeafAncestors(ancestors);
                return ancestors.size();
        }

        /**
         * <p>
         * Returns the number of trie nodes that are leafs.
         * </p>
         * 
         * @return number of leafs in the trie
         */
        public int getNumLeafs() {
                return rootNode.getNumLeafs();
        }
}