题目
前缀树
class AutocompleteSystem {
private StringBuilder haveInput;
private List<Pair<String, Integer>> ans;
public AutocompleteSystem(String[] sentences, int[] times) {
haveInput = new StringBuilder();
ans = new ArrayList<>();
root = new TrieNode();
for (int i = 0; i < sentences.length; i++) {
insert(sentences[i], times[i]);
}
}
public List<String> input(char c) {
ans.clear();
List<String> res = new ArrayList<>();
if (c == '#') {
insert(haveInput.toString(), 1);
haveInput.setLength(0);
} else {
haveInput.append(c);
TrieNode node = searchNode(haveInput.toString());
search(node);
Collections.sort(ans, (o1, o2) -> {
if (o1.getValue() == o2.getValue())
return o1.getKey().compareTo(o2.getKey());
return o2.getValue().compareTo(o1.getValue());
});
for (int i = 0; i < Math.min(3, ans.size()); i++) {
res.add(ans.get(i).getKey());
}
}
return res;
}
private TrieNode root;
class TrieNode {
String sentence;
int freq;
Map<Character, TrieNode> children;
TrieNode() {
sentence = null;
freq = 0;
children = new HashMap<>();
}
}
public void insert(String sentence, int time) {
TrieNode cur = root;
for (char c : sentence.toCharArray()) {
TrieNode next = cur.children.computeIfAbsent(c, k -> new TrieNode());
cur = next;
}
cur.freq += time;
cur.sentence = sentence;
}
private TrieNode searchNode(String s) {
TrieNode cur = root;
for (char c : s.toCharArray()) {
if (cur.children.containsKey(c)) {
cur = cur.children.get(c);
} else {
return null;
}
}
return cur;
}
private void search(TrieNode node) {
if (node == null) return;
if (node.sentence != null) {
ans.add(new Pair(node.sentence, node.freq));
}
for (TrieNode child : node.children.values()) {
search(child);
}
}
}
/**
* Your AutocompleteSystem object will be instantiated and called as such:
* AutocompleteSystem obj = new AutocompleteSystem(sentences, times);
* List<String> param_1 = obj.input(c);
*/