Java Trie + PriorityQueue and further thoughts


  • 0
    M
        public class AutocompleteSystem {
    	private class Node {
    		Node[] next = new Node[128];
    		boolean isWord = false;
    		Map<String, Integer> cnt  = new HashMap<>();
    		public Node(){
    			
    		}
    	}
    	private class Trie{
    		Node root = new Node();
    		public Trie(String[] sentences, int[] times) {
    			for (int i = 0; i < sentences.length; ++i) {
    				add(sentences[i], times[i]);
    			}
    		}
    		public void add(String s, int time) {
    			Node cur = root;
    			for (char c : s.toCharArray()) {
    				if (cur.next[c] == null) cur.next[c] = new Node();
    				cur = cur.next[c];
    				cur.cnt.put(s, cur.cnt.getOrDefault(s, 0) + time);
    			}
    			cur.isWord = true;
    		}
    	}
    	private class Pair implements Comparable<Pair>{
    		String word;
    		int cnt;
    		public Pair(String word, int cnt) {
    			this.word = word;
    			this.cnt = cnt;
    		}
    		@Override
    		public int compareTo(Pair b) {
    			if (cnt != b.cnt) return cnt - b.cnt;
    			return b.word.compareTo(word);
    		}
    	}
    	private Trie trie;
    	private StringBuilder typed;
    	private Node node;
        public AutocompleteSystem(String[] sentences, int[] times) {
            trie = new Trie(sentences, times);
    		node = trie.root;
    		typed = new StringBuilder();
        }
        
        public List<String> input(char c) {
            if (c == '#') {
    			trie.add(typed.toString(), 1);
    			typed = new StringBuilder();
    			node = trie.root;
    			return new ArrayList<String>();
    		}
    		if (node.next[c] == null) {
    			node.next[c] = new Node();			
    		}
    		node = node.next[c];
    		typed.append(c);
    		Queue<Pair> minHeap = new PriorityQueue<>(3);
    		for (String key : node.cnt.keySet()) {
    			Pair pair = new Pair(key, node.cnt.get(key));
    			if (minHeap.size() < 3) minHeap.offer(pair);
    			else {
    				if (minHeap.peek().compareTo(pair) < 0) {
    					minHeap.poll();
    					minHeap.offer(pair);
    				}
    			}
    		}
    		List<String> list = new LinkedList<String>();
    		while (!minHeap.isEmpty()) {
    			list.add(0, minHeap.poll().word);
    		}
    		return list;
        }
    }
    

    Actually there must be better design, like we can build trie with nodes that contain HashHeap which supports key increment in logarithmic time complexity. However I think the implementation of HashHeap is too complex to fit in an interview.


Log in to reply
 

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