链接:
1268. 搜索推荐系统
题解:
class Solution {
public:
struct Trie {
Trie() {
end = false;
next.resize(26, nullptr);
}
bool end;
std::set<std::string> words;
std::vector<Trie*> next;
};
void insert_trie(const std::string& word) {
Trie* cur = _root;
for (auto ch : word) {
if (!cur->next[ch-'a']) {
cur->next[ch-'a'] = new (std::nothrow)Trie;;
}
cur = cur->next[ch-'a'];
cur->words.insert(word);
}
cur->end = true;
cur->words.insert(word);
}
std::vector<std::string> find_words(const std::string& prefix) {
Trie* cur = _root;
for (int i = 0; i < prefix.size(); ++i) {
if (!cur->next[prefix[i]-'a']) {
return {};
}
cur = cur->next[prefix[i]-'a'];
}
if (cur->words.size() <= 3) {
std::vector<std::string> result(cur->words.begin(), cur->words.end());
return result;
} else {
std::vector<std::string> result;
auto i = 0;
for (auto ite = cur->words.begin(); i < 3; ++i, ++ite) {
result.push_back(*ite);
}
return result;
}
return {};
}
public:
vector<vector<string>> suggestedProducts(vector<string>& products, string searchWord) {
_root = new (std::nothrow)Trie;
for (auto& product : products) {
insert_trie(product);
}
std::vector<std::vector<std::string>> result;
result.reserve(searchWord.size());
for (int i = 0; i < searchWord.size(); ++i) {
result.push_back(find_words(searchWord.substr(0, i+1)));
}
return result;
}
private:
Trie* _root;
};