TrieTree

TrieTree.java

package com.shenbian.service.personchange.impl.utils;

import java.util.*;

/**
 * Trie Tree
 * @author JobsDong
 */
public class TrieTree {

    private TrieNode root = new TrieNode('A');

    /**
     * insert word with some property
     * @param term word
     * @param properties some property wanted
     */
    public void insert(String term, Map<String, Object> properties) {
        TrieNode trieNode = root;
        for (char character : term.toCharArray()) {

            boolean isFind = false;
            for (TrieNode child : trieNode.children) {
                if (child.character == character) {
                    trieNode = child;
                    isFind = true;
                    break;
                }
            }
            if (!isFind) {
                TrieNode newChild = new TrieNode(character);
                trieNode.addChild(newChild);
                trieNode = newChild;
            }
        }

        trieNode.isTerm = true;
        trieNode.setProperties(properties);
    }

    /**
     * find term in trie tree
     * @param term
     * @return
     */
    public TrieNode searchTerm(String term) {
        TrieNode trieNode = root;
        for (char character : term.toCharArray()) {
            boolean isFind = false;
            for (TrieNode child : trieNode.children) {
                if (child.character == character) {
                    trieNode = child;
                    isFind = true;
                    break;
                }
            }
            if (!isFind) {
                trieNode = null;
                break;
            }
        }

        if (trieNode != null && trieNode.isTerm) {
            return trieNode;
        } else {
            return null;
        }
    }

    /**
     * find terms in a text
     * @param text
     * @return
     */
    public List<Term> searchText(String text) {
        List<Term> terms = new ArrayList<Term>();

        for (int termStart = 0; termStart < text.length(); ) {
            int termLen = 0;
            TrieNode trieNode = root;
            int maxTermLen = 0;
            TrieNode maxTrieNode = root;

            while (termStart + termLen < text.length()) {
                char character = text.charAt(termStart + termLen);
                boolean isFind = false;
                for (TrieNode child : trieNode.children) {
                    if (child.character == character) {
                        trieNode = child;
                        termLen++;
                        if (trieNode.isTerm) {
                            maxTermLen = termLen;
                            maxTrieNode = trieNode;
                        }
                        isFind = true;
                        break;
                    }
                }
                if (!isFind) {
                    break;
                }
            }

            if (maxTrieNode != root) {
                terms.add(new Term(text.substring(termStart, termStart+maxTermLen),
                        termStart, maxTrieNode.getProperties()));
            }

            termStart += maxTermLen + 1;
        }

        return terms;
    }

    /**
     * @return
     */
    public String traverse() {
        Deque<TrieNode> deque = new ArrayDeque<TrieNode>();
        deque.addLast(root);

        StringBuilder sb = new StringBuilder();
        while (!deque.isEmpty()) {
            TrieNode node = deque.removeFirst();
            for (TrieNode child : node.children) {
                sb.append(String.format("<node:%s %d>", child.character, child.children.size()));
                deque.addLast(child);
            }
        }

        return sb.toString();
    }

    public static class TrieNode {
        private char character;
        private List<TrieNode> children = new ArrayList<TrieNode>();
        private boolean isTerm = false;

        private Map<String, Object> properties = new HashMap<String, Object>();

        TrieNode(char character) {
            this.character = character;
        }

        void addChild(TrieNode child) {
            children.add(child);
        }

        void setProperties(Map<String, Object> properties) {
            if (properties == null) {
                properties = new HashMap<String, Object>();
            }
            this.properties = properties;
        }

        Object getProperty(String key) {
            return properties.get(key);
        }

        Map<String, Object> getProperties() {
            return properties;
        }
    }

    public static class Term {
        private int start;
        private String term;
        private Map<String, Object> properties = new HashMap<String, Object>();

        public Term(String term, int start, Map<String, Object> properties) {
            this.start = start;
            this.term = term;
            this.properties = properties;

        }

        public String getTerm() {
            return term;
        }

        public int getStart() {
            return start;
        }

        public void setProperties(Map<String, Object> properties) {
            if (properties == null) {
                properties = new HashMap<String, Object>();
            }
            this.properties = properties;
        }

        public Object getProperty(String key) {
            return properties.get(key);
        }
    }

    public static void main(String[] args) {
        TrieTree trieTree = new TrieTree();
        trieTree.insert("word", new HashMap<String, Object>());
        trieTree.insert("worry", new HashMap<String, Object>());
        //search term
        TrieNode trieNode = trieTree.searchTerm("word");
        assert (trieNode != null);
        trieNode = trieTree.searchTerm("hello");
        assert (trieNode == null);

        //search terms in text
        List<Term> terms = trieTree.searchText("I'm worry about you. In a word, i love you.");
        assert (terms.size() == 2);
        assert (terms.get(0).getTerm().equals("worry"));
        assert (terms.get(1).getTerm().equals("word"));

        System.out.println(trieTree.traverse());
    }
}