Trie: Insert, Search and Prefix Count

/*
 * Problem : Process commands line by line.
 * Input   : Several lines of commands.
 * Output  : One answer per command.
 * Approach: A tree of letters.
 * Time    : O(L) per command
 * Space   : O(total letters)
 * Java    : 8 or newer.
 */
import java.util.Scanner;

public class Main {
        // Step 1: A trie node has up to 26 children (one for each letter). passCount = how many words go through this node.
        static class Node {
            Node[] children = new Node[26];
            int passCount = 0;
            boolean isWord = false;
        }

        static class Trie {
            private final Node root = new Node();

            // Step 2: To insert a word, walk down the letters. Create missing nodes. Count each node we pass.
            void insert(String word) {
                Node node = root;
                for (int i = 0; i < word.length(); i++) {
                    int index = word.charAt(i) - 'a';
                    if (node.children[index] == null) {
                        node.children[index] = new Node();
                    }
                    node = node.children[index];
                    node.passCount++;
                }
                node.isWord = true;
            }

            // Step 3: Walk down the letters of the prefix. Return null if a letter is missing.
            private Node walk(String text) {
                Node node = root;
                for (int i = 0; i < text.length(); i++) {
                    node = node.children[text.charAt(i) - 'a'];
                    if (node == null) {
                        return null;
                    }
                }
                return node;
            }

            int countPrefix(String prefix) {
                Node node = walk(prefix);
                return node == null ? 0 : node.passCount;
            }

            boolean search(String word) {
                Node node = walk(word);
                return node != null && node.isWord;
            }
        }

    // This method solves the task. It receives the input and returns the answer.
    static String[] solve(String[] lines) {
        Trie trie = new Trie();
        String[] output = new String[lines.length];
        for (int i = 0; i < lines.length; i++) {
            String[] parts = lines[i].split(" ");

            // Step 4: Run the command.
            if (parts[0].equals("insert")) {
                trie.insert(parts[1]);
                output[i] = "ok";
            } else if (parts[0].equals("count")) {
                output[i] = String.valueOf(trie.countPrefix(parts[1]));
            } else {
                output[i] = String.valueOf(trie.search(parts[1]));
            }
        }
        return output;
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        java.util.List<String> all = new java.util.ArrayList<>();
        while (scanner.hasNextLine()) {
            String line = scanner.nextLine().trim();
            if (!line.isEmpty()) {
                all.add(line);
            }
        }
        String[] lines = all.toArray(new String[0]);
        System.out.println(String.join("\n", solve(lines)));
    }
}