Trie字符串统计 For Java

原题链接：Trie字符串统计

Trie对象版

ps: 这个稍微快点

import java.util.Scanner;
import java.io.*;

public class Main {
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        PrintWriter out = new PrintWriter(System.out);
        Trie trie = new Trie();

        int n = Integer.parseInt(in.nextLine());
        while (n -- > 0) {
            String[] line = in.nextLine().split(" ");
            if (line[0].equals("I")) trie.insert(line[1]);
            else out.println(trie.query(line[1]));
        }
        out.flush();
    }

    static class Trie {
        private Trie[] trie;
        private int count;

        public Trie() {
            trie = new Trie[26];
        }

        public void insert(String str) {
            Trie root = this;
            for (char c : str.toCharArray()) {
                int u = c - 'a';
                if (root.trie[u] == null) root.trie[u] = new Trie();
                root = root.trie[u];
            }
            root.count++;
        }

        public int query(String str) {
            Trie root = this;
            for (char c : str.toCharArray()) {
                int u = c - 'a';
                if (root.trie[u] == null) return 0;
                root = root.trie[u];
            }
            return root.count;
        }
    }
}

Trie数组版

import java.util.*;
import java.io.*;

public class Main {
    private static int[][] son;
    private static int[] cnt;
    //下标为0的点，既是空节点，又是根节点
    private static int idx;
    public static void main(String[] args) {
        Scanner in = new Scanner(System.in);
        PrintWriter out = new PrintWriter(System.out);
        son = new int[100010][26];
        cnt = new int[100010];

        int n = Integer.parseInt(in.nextLine());
        while (n -- > 0) {
            String[] line = in.nextLine().split(" ");
            if (line[0].equals("I")) insert(line[1]);
            else out.println(query(line[1]));
        }
        out.flush();
    }

    public static void insert(String str) {
        int p = 0;
        for (char c : str.toCharArray()) {
            int u = c - 'a';
            if (son[p][u] == 0) son[p][u] = ++idx;
            p = son[p][u];
        }
        cnt[p]++;
    }

    public static int query(String str) {
        int p = 0;
        for (char c : str.toCharArray()) {
            int u = c - 'a';
            if (son[p][u] == 0) return 0;
            p = son[p][u];
        }
        return cnt[p];
    }
}