前缀树工具类源码

前缀树工具类源码及测试

public class Trie {

private TrieNode root;
public Trie(){
    root = new TrieNode();
    root.var = ' ';
}

/**
 * 插入trie树
 * @param word
 */
public void insert(String word){
    TrieNode ws = root;
    for(int i = 0; i < word.length(); i++){
        char c = word.charAt(i);
        if(!ws.children.keySet().contains(c)){
            ws.children.put(c,new TrieNode(c));
        }
        ws = ws.children.get(c);
    }
    ws.isWord = true;
}


/**
 * 查询trie树
 * @param prefix
 * @return
 */
public List<String> startWith(String prefix){
    List<String> match = new ArrayList<>();
    TrieNode ws = root;
    for(int i = 0; i < prefix.length(); i++){
        char c = prefix.charAt(i);
        if(!ws.children.keySet().contains(c)) return match;
        ws = ws.children.get(c);
        if(!ws.containLongTail){
            for (char cc : ws.children.keySet()){
                match.add(prefix+cc);
            }
        }
    }
    return match;
}

public static void main(String[] args) {
    Trie t = new Trie();
    t.insert("阿辉");
    List<String> ret = t.startWith("阿");
    System.out.println(ret);
}

}