编程-单词替换解法--前缀树

struct TrieTree {
    bool endflag;
    map<char, TrieTree*> son;
    TrieTree() : endflag(false) {}
};

class Solution {

public:
    TrieTree* root;

    Solution() {
        root = new TrieTree();
    }

    string replaceWords(vector<string>& dict, string sentence) {

        stringstream ss(sentence);
        string ret;
        string init;

        if(sentence.empty())
            return ret;

        TrieTree* node;
        for (auto w : dict) {
            node = root;
            for (auto x : w) {
                if ((node->son).count(x) == 0) {
                    node->son[x] = new TrieTree();
                }
                node = node->son[x];
            }
            node->endflag = true;
        }


        while(ss>>init) {
            node = root;
            string s1;
            for (auto x : init) {
                if ((node->son).count(x) != 0) {
                    s1 += x;
                    node = node->son[x];
                } else {
                    break;
                }
                if (node->endflag)
                    break;
            }
            ret += node->endflag ? s1 + ' ' : init + ' ';
        }
        if (!ret.empty()) {
            ret.pop_back();
        }
        return ret;
    }
};

int main()
{
    vector<string> dictionary = {"a", "aa", "aaa", "aaaa"};
    string sentence = "a aa a aaaa aaa aaa aaa aaaaaa bbb baba ababa";
    Solution s;
    cout << s.replaceWords(dictionary,sentence) << endl;

    return 0;
}