1268. 搜索推荐系统

 

链接：

1268. 搜索推荐系统

题解：

class Solution {
public:
struct Trie {
    Trie() {
        end = false;
        next.resize(26, nullptr);
    }
    bool end;
    std::set words;
    std::vector 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 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 result(cur->words.begin(), cur->words.end());
            return result;
        } else {
            std::vector result;
            auto i = 0;
            for (auto ite = cur->words.begin(); i  < 3; ++i, ++ite) {
                result.push_back(*ite);
            }
            return result;
        }
        return {};
    }
public:
    vector> suggestedProducts(vector& products, string searchWord) {
        _root = new (std::nothrow)Trie;
        for (auto& product : products) {
            insert_trie(product);
        }
        std::vector> 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;
};