Normal Trie (with parent pointer) beats 90%+


  • 1
    N

    Trie with children and a parent pointer.
    When adding a sentence, get to the bottom node and bubble up all the way up to root node, updating the top3 list (only used list here since there's only top 3. I tried to use TreeSet but failed the last case. Looks like duplicate nodes are added in TreeSet. So much thanks if anyone can help me debug.. )

    import java.util.*;
    
    public class AutocompleteSystem {
        private TrieNode rootNode;
        private TrieNode curNode;
        private StringBuilder input;
    
        private class TrieNode implements Comparable<TrieNode>{
            String sentence;
            int count;
            TrieNode parent;
            TrieNode[] children;
            List<TrieNode> top3;
    
            public TrieNode(TrieNode parent) {
                this.parent = parent;
                children = new TrieNode[128];
                top3 = new ArrayList<TrieNode>();
            }
            public int compareTo(TrieNode node) {
                if (this.count == node.count) {
                    return this.sentence.compareTo(node.sentence);
                }
                return node.count - this.count;
            }
        }
    
        private boolean checkEqual(TrieNode node1, TrieNode node2) {
            return node1.sentence.equals(node2.sentence);
        }
    
        private void addSentence(String sentence, int count) {
            TrieNode node = rootNode;
    
            // get to the leave node.
            for (char c : sentence.toCharArray()) {
                if (node.children[c] == null) {
                    node.children[c] = new TrieNode(node);   // set parent as current node
                }
                node = node.children[c];
            }
            node.sentence = sentence;
            node.count = node.count + count;
    
            // pop up from leave node all the way to root node.  update top3
            TrieNode cur = node;
            while (cur.parent != null) {
                if (cur.top3.size() == 0) {
                    cur.top3.add(node);
                }
                else {
                    boolean contains = false;
                    for (TrieNode n : cur.top3) {
                        if (checkEqual(n, node)) {
                            contains = true;
                            break;
                        }
                    }
    
                    if (!contains) {
                        cur.top3.add(node);
                    }
                    Collections.sort(cur.top3);
                    if (cur.top3.size() > 3) {
                        cur.top3.remove(cur.top3.size() - 1);
                    }
                }
    
                cur = cur.parent;
            }
    
    
        }
    
        // assume curNode never be null
        private List<String> inputHelper(char c){
            input.append(c);
    
            if (curNode == null || curNode.children[c] == null) {
                curNode = null;
                return new ArrayList<String>();
            }
    
            List<String> res = new ArrayList<>();
            TrieNode next = curNode.children[c];
            for (TrieNode node : next.top3) {
                res.add(node.sentence);
            }
    
            curNode = next;
            return res;
        }
    
        public AutocompleteSystem(String[] sentences, int[] times) {
            rootNode = new TrieNode(null);
            curNode = rootNode;
            input = new StringBuilder();
    
            for (int i=0; i<sentences.length; i++) {
                addSentence(sentences[i], times[i]);
            }
    
        }
    
        public List<String> input(char c) {
            if (c == '#') {
                addSentence(input.toString(), 1);
    
                input = new StringBuilder();
                curNode = rootNode;
                return new ArrayList<String>();
            }
            else {
                return inputHelper(c);
            }
        }
    
    
    }
    

    If anyone has time, there's another version using TreeSet() in the trie nodes to keep track of top3. However it won't pass the last test case... please help..

    I figure it should be something about TreeSet.contains()..

    import java.util.*;
    
    public class AutocompleteSystem {
        private TrieNode rootNode;
        private TrieNode curNode;
        private StringBuilder input;
    
        public class TrieNode implements Comparable<TrieNode>{
            String sentence;
            int count;
            TrieNode parent;
            TrieNode[] children;
            TreeSet<TrieNode> top3;
    
            public TrieNode(TrieNode parent) {
                this.parent = parent;
                children = new TrieNode[128];
                top3 = new TreeSet<TrieNode>();
            }
    
            public int compareTo(TrieNode node) {
                if (node.count == this.count) {
                    return node.sentence.compareTo(this.sentence);
                }
                return this.count - node.count;
            }
    
    
        }
    
        public void addSentence(String sentence, int count) {
            TrieNode node = rootNode;
    
            for (char c : sentence.toCharArray()) {
                if (node.children[c] == null) {
                    node.children[c] = new TrieNode(node);   // set parent as current node
                }
                node = node.children[c];
            }
    
            node.sentence = sentence;
            node.count = node.count + count;   // could end up finding existing node. add the count
    
            TrieNode cur = node;
            while(cur.parent != null) {
                TreeSet set = cur.top3;
                if (set.contains(node)) {
                    set.remove(node);
                }
                set.add(node);
    
                if (set.size() > 3) {
                    set.pollFirst();
                }
    
                cur = cur.parent;
            }
    
        }
    
        // assume curNode never be null
        private List<String> inputHelper(char c){
            input.append(c);
    
    
            if (curNode == null || curNode.children[c] == null) {
                curNode = null;
                return new ArrayList<String>();
            }
    
            List<String> res = new ArrayList<>();
            TrieNode next = curNode.children[c];
    
            Iterator<TrieNode> it = next.top3.descendingIterator();
            while (it.hasNext()) {
                res.add(it.next().sentence);
            }
    
            curNode = next;
            return res;
        }
    
        public AutocompleteSystem(String[] sentences, int[] times) {
            rootNode = new TrieNode(null);
            curNode = rootNode;
            input = new StringBuilder();
    
            for (int i=0; i<sentences.length; i++) {
                addSentence(sentences[i], times[i]);
            }
    
        }
    
        public List<String> input(char c) {
            if (c == '#') {
                addSentence(input.toString(), 1);
    
                input = new StringBuilder();
                curNode = rootNode;
                return new ArrayList<String>();
            }
            else {
                return inputHelper(c);
            }
        }
    
    
    }
    

Log in to reply
 

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