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);
            }
        }
    
    
    }
    

  • 0
    K

    TreeSet is not like HashSet. It uses compareTo() values to ensure the uniqueness.
    When you use node.count = node.count + count to change the count value, the contains() won't work any more for this node.


Log in to reply
 

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