Java Trie solution which doesn't traverse all children when calling input()


  • 0
    A
    public class AutocompleteSystem {
        
        class TrieNode {
            String s;
            int t;
            List<TrieNode> data;
            TrieNode[] children;
            
            TrieNode() {
                data = new ArrayList<TrieNode>();
                children = new TrieNode[27];
            }
        }
        
        private int idx(char c) {
            return c == ' ' ? 26 : c - 'a';
        }
        
        private void putIn(List<TrieNode> data, TrieNode node) {
            boolean isFound = false;
            for (TrieNode d : data) {
                if (d == node) {
                    isFound = true;
                    break;
                }
            }
            
            if (!isFound) {
                data.add(node);
            }
            
            Collections.sort(data, (a, b) -> a.t == b.t ? a.s.compareTo(b.s) : b.t - a.t);
            
            if (data.size() > 3) {
                data.remove(data.size() - 1);
            }
        }
        
        private void insert(TrieNode root, String s, int t) {
            TrieNode node = root;
            for (int i = 0; i < s.length(); i++) {
                int j = idx(s.charAt(i));
                if (node.children[j] == null) {
                    node.children[j] = new TrieNode();
                }
                node = node.children[j];
            }
            node.s = s;
            node.t += t;
            
            TrieNode temp = root;
            for (int i = 0; i < s.length(); i++) {
                int j = idx(s.charAt(i));
                temp = temp.children[j];
                putIn(temp.data, node);
            }
        }
        
        private TrieNode root;
        private TrieNode curr;
        private String currStr;
        
        public AutocompleteSystem(String[] sentences, int[] times) {
            curr = root = new TrieNode();
            currStr = "";
            for (int i = 0; i < sentences.length; i++) {
                insert(root, sentences[i], times[i]);
            }
        }
        
        public List<String> input(char c) {
            List<String> ans = new ArrayList<>();
            if (c == '#') {
                insert(root, currStr, 1);
                currStr = "";
                curr = root;
                return ans;
            }
            
            currStr += c;
            
            int j = idx(c);
            if (curr.children[j] == null) {
                curr.children[j] = new TrieNode();
            }
            curr = curr.children[j];
            
            for (TrieNode node : curr.data) {
                ans.add(node.s);
            }
            return ans;
        }
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.