Trie 树

实现一个trie树，快速返回包含某个前缀的单词数量。

http://hihocoder.com/problemset/problem/1014

import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintStream;
import java.util.HashMap;
import java.util.Map;
import java.util.Scanner;

public class Main {

    public static void main(String[] args) throws FileNotFoundException {
        TrieNode root = new TrieNode('1');

        // Scanner in = new Scanner(new File("A.in"));
        // PrintStream out = new PrintStream(new File("A.out"));
        Scanner in = new Scanner(System.in);
        PrintStream out = new PrintStream(System.out);

        int n = in.nextInt();
        for (int i = 0; i < n; i++) {
            String str = in.next();
            root.insert(str);
        }

        int m = in.nextInt();
        for (int i = 0; i < m; i++) {
            String str = in.next();
            out.println(root.getCount(str));
        }

    }
}

class TrieNode {
    Character curChar;
    Map<Character, TrieNode> children = new HashMap<Character, TrieNode>();
    int count = 1;

    public TrieNode(char ch) {
        curChar = ch;
    }

    public void insert(String str) {
        if (str.isEmpty())
            return;
        char cur = str.charAt(0);
        if (children.containsKey(cur)) {
            TrieNode child = children.get(cur);
            child.count++;
            child.insert(str.substring(1));
        } else {
            TrieNode newNode = new TrieNode(cur);
            children.put(cur, newNode);
            newNode.insert(str.substring(1));
        }
    }

    public int getCount(String str) {
        if (str == null)
            return 0;
        if (str.isEmpty()) {
            return count;
        }
        char ch = str.charAt(0);
        if (!children.containsKey(ch)) {
            return 0;
        } else
            return children.get(ch).getCount(str.substring(1));
    }
}