Search Autocomplete

1class AutocompleteDfs {2    static class Node {3        final Node[] next = new Node[26]; // 'a'..'z'4        int weight; // > 0 only where a word ends5    }6 7    private final Node root = new Node();8 9    public void addWord(String word, int weight) {10        Node node = root;11        for (char ch : word.toCharArray()) {12            int i = ch - 'a';13            if (node.next[i] == null) node.next[i] = new Node();14            node = node.next[i];15        }16        node.weight += weight; // the same word again: the weights add up17    }18 19    public List<String> suggest(String prefix, int k) {20        Node node = root;21        for (char ch : prefix.toCharArray()) { // walk down the prefix: O(p)22            node = node.next[ch - 'a'];23            if (node == null) return new ArrayList<>();24        }25        // Collect every word under the prefix, keep the best k in a min-heap (worst of the best on top).26        Comparator<Map.Entry<String, Integer>> better = (a, b) ->27                !a.getValue().equals(b.getValue()) ? Integer.compare(b.getValue(), a.getValue()) : a.getKey().compareTo(b.getKey());28        PriorityQueue<Map.Entry<String, Integer>> best = new PriorityQueue<>(better.reversed());29        collect(node, new StringBuilder(prefix), best, k);30        List<String> result = new ArrayList<>();31        while (!best.isEmpty()) result.add(best.poll().getKey()); // worst first…32        Collections.reverse(result); // …so flip it33        return result;34    }35 36    /** Visits the whole subtree: O(size of the subtree · log k). Recursion depth = word length. */37    private void collect(Node node, StringBuilder path, PriorityQueue<Map.Entry<String, Integer>> best, int k) {38        if (node.weight > 0) {39            best.add(Map.entry(path.toString(), node.weight));40            if (best.size() > k) best.poll(); // drop the worst41        }42        for (int i = 0; i < 26; i++) {43            if (node.next[i] == null) continue;44            path.append((char) ('a' + i));45            collect(node.next[i], path, best, k);46            path.deleteCharAt(path.length() - 1);47        }48    }49}

1class Autocomplete {2    static final int MAX_K = 10; // the spec promises k <= 10, so each node keeps its best 103 4    static class Node {5        final Node[] next = new Node[26];6        final List<String> top = new ArrayList<>(); // best words under this node, best first, at most MAX_K7    }8 9    private final Node root = new Node();10    private final Map<String, Integer> weights = new HashMap<>(); // word → total weight11 12    public void addWord(String word, int weight) {13        weights.merge(word, weight, Integer::sum);14        Node node = root;15        update(node.top, word); // the root answers the empty prefix ""16        for (char ch : word.toCharArray()) {17            int i = ch - 'a';18            if (node.next[i] == null) node.next[i] = new Node();19            node = node.next[i];20            update(node.top, word); // every prefix of the word may now rank it differently21        }22    }23 24    /**25     * Puts word in its right place in one node's list, then trims it to MAX_K. O(MAX_K).26     * Safe because weights only grow: a word that fell off a list can come back only when its own weight27     * grows, and that call walks this same path.28     */29    private void update(List<String> top, String word) {30        top.remove(word);31        int w = weights.get(word);32        int at = 0;33        while (at < top.size() && isBetter(top.get(at), weights.get(top.get(at)), word, w)) at++;34        top.add(at, word);35        if (top.size() > MAX_K) top.remove(top.size() - 1);36    }37 38    private static boolean isBetter(String a, int wa, String b, int wb) {39        return wa != wb ? wa > wb : a.compareTo(b) < 0;40    }41 42    public List<String> suggest(String prefix, int k) {43        Node node = root;44        for (char ch : prefix.toCharArray()) { // O(p), then the answer is already waiting45            node = node.next[ch - 'a'];46            if (node == null) return new ArrayList<>();47        }48        return new ArrayList<>(node.top.subList(0, Math.min(k, node.top.size())));49    }50}