1268. Search Suggestions System

Question

You are given an array of strings products and a string searchWord.

Design a system that suggests at most three product names from products after each character of searchWord is typed. Suggested products should have common prefix with searchWord. If there are more than three products with a common prefix return the three lexicographically minimums products.

Return a list of lists of the suggested products after each character of searchWord is typed.

Solution

字典树+优先级队列。
在TrieNode节点中保存PriorityQueue，以记录该节点下的三个单词。
注意PriorityQueue中的String需要按倒序排列，以保证维护时可以直接poll掉字符顺序较后的单词。

add()方法

创建并添加TrieNode节点，同时在移动当前节点位置时，需要维护PriorityQueue，当长度大于3时，则挤出字典顺序排列较后的单词。

search()方法

在查找前缀时，将每个TrieNode节点上的PriorityQueue倒序加入List。

Code

class Solution {
    class TrieNode{
        boolean end;
        TrieNode[] children = new TrieNode[26];
        PriorityQueue<String> q = new PriorityQueue<>((a, b) -> b.compareTo(a));
    }
    
    TrieNode root;
    
    public List<List<String>> suggestedProducts(String[] products, String searchWord) {
        List<List<String>> ret = new ArrayList<>();
        root = new TrieNode();
        for(String word : products){
            add(word);
        }
        return search(searchWord);
    }
    
    private void add(String word){
        TrieNode curr = root;
        for(char c : word.toCharArray()){
            int i = c - 'a';
            if(curr.children[i] == null) curr.children[i] = new TrieNode();
            curr = curr.children[i];
            curr.q.add(word);
            if(curr.q.size() > 3) curr.q.poll();
        }
        curr.end = true;
    }
    
    private List<List<String>> search(String prefix){
        List<List<String>> ret = new ArrayList<>();
        TrieNode curr = root;
        for(char c : prefix.toCharArray()){
            int i = c - 'a';
            if(curr.children[i] == null) curr.children[i] = new TrieNode();
            curr = curr.children[i];
            List<String> temp = new ArrayList<>();
            while(!curr.q.isEmpty()){
                temp.add(curr.q.poll());
            }
            Collections.reverse(temp);
            ret.add(temp);
        }
        return ret;
    }
}