Java Trie Solution Avoiding traversal from root and subtree


  • 1
    A

    Trading space for time. Main improvements are :

    1. avoiding traversal from root for each input invocation by keeping track of current node and current sentence.
    2. Using a string buffer instead of string.
    3. Avoiding subtree traversal to get list of matching sentences . Each node maintains a list of sentence indices.
    class AutocompleteSystem {
        
        private static class TrieNode {
            private static final int MAX_LINKS = 27;
    
            private TrieNode[] links;
    
            private boolean end;
    
            private Set<Integer> sentenceIndices;
    
            public TrieNode() {
                links = new TrieNode[MAX_LINKS];
                sentenceIndices = new HashSet<>();
            }
    
            public boolean containsKey(final char ch) {
                return links[charIndex(ch)] != null;
            }
    
            public TrieNode get(final char ch) {
                return links[charIndex(ch)];
            }
    
            public void put(final char ch, final TrieNode node) {
                links[charIndex(ch)] = node;
            }
    
            public boolean isEnd() {
                return end;
            }
    
            public void setEnd() {
                this.end = true;
            }
    
            public void addSentenceIndex(final int index) {
                if (!sentenceIndices.contains(index)) {
                    sentenceIndices.add(index);
                }
            }
    
            private int charIndex(final char ch) {
                return ch != ' ' ? ch - 'a' : links.length-1;//links[26] will represent space character
            }
        }
    
        private List<String> sentences;
        private Map<String, Integer> sentenceCountMap;
    
        private TrieNode root;
        private StringBuffer buffer;
        private TrieNode currentNode;
    
        public AutocompleteSystem(final String[] sentences, final int[] times) {
            this.sentences = new ArrayList<>(Arrays.asList(sentences));
            this.sentenceCountMap = new HashMap<>();
            for (int i = 0; i < sentences.length; i++) {
                sentenceCountMap.put(sentences[i], times[i]);
            }
    
            this.root = new TrieNode();
            this.buffer = new StringBuffer();
            this.currentNode = root;
    
            for (int i = 0; i < sentences.length; i++) {
                insert(i);
            }
        }
    
        public List<String> input(final char c) {
            if (c == '#') {
                String sentence = this.buffer.toString();
                this.buffer.setLength(0);
                this.currentNode = root;
    
                if (!sentenceCountMap.containsKey(sentence)) {
                    sentences.add(sentence);
                    insert(sentences.size()-1);
                }
                sentenceCountMap.put(sentence, sentenceCountMap.getOrDefault(sentence, 0) + 1);
                return Collections.emptyList();
            } else {
                this.buffer.append(c);
    
                if (currentNode == null) {// no more matching is possible
                    return Collections.emptyList();
                } else if (!currentNode.containsKey(c)) {// no more matching is possible
                    currentNode = null;
                    return Collections.emptyList();
                } else {
                    currentNode = currentNode.get(c);
    
                    List<Integer> candidateIndices = new ArrayList<>();
                    candidateIndices.addAll(currentNode.sentenceIndices);
                    Collections.sort(candidateIndices, (index1, index2) -> {
                        String sentence1 = sentences.get(index1);
                        String sentence2 = sentences.get(index2);
                        int count1 = sentenceCountMap.get(sentence1);
                        int count2 = sentenceCountMap.get(sentence2);
                        return count1 == count2 ? sentence1.compareTo(sentence2) : Integer.compare(count2, count1);
                    });
    
                    final List<Integer> resultIndices = candidateIndices.subList(0, Math.min(3, candidateIndices.size()));
                    return resultIndices.stream().map(index -> this.sentences.get(index)).collect(Collectors.toList());
                }
            }
        }
    
        private void insert(final int sentenceIndex) {
            TrieNode current = root;
            String sentence = sentences.get(sentenceIndex);
            for (char ch : sentence.toCharArray()) {
                if (!current.containsKey(ch)) {
                    current.put(ch, new TrieNode());
                }
                current = current.get(ch);
                current.addSentenceIndex(sentenceIndex);
            }
            current.setEnd();
        }
    }
    

Log in to reply
 

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