(trie) 648. 单词替换

648. 单词替换

const int N = 1e6+10;
int son[N][26], idx = 0;
int cnt[N];

class Solution {
public:
    void insert(string s){
        int p = 0;
        for(auto c: s){
            int u = c - 'a';
            if(!son[p][u]) son[p][u] = ++idx;
            p = son[p][u];
        }
        cnt[p] = 1;
    }
    string query(string s){
        int p = 0;
        string res = "";
        for(auto c: s){
            int u = c - 'a';
            if(!son[p][u]) return s;
            p = son[p][u];
            res += c;
            if(cnt[p]) return res;
        }
        return s;
    }
    string replaceWords(vector<string>& dictionary, string sentence) {
        memset(son, 0, sizeof son);
        memset(cnt, 0, sizeof cnt);
        for(auto s: dictionary){
            insert(s);
        }
        stringstream ssin(sentence);
        string res = "", w;
        while(ssin >> w){
            string curs = query(w);
            res += curs + " ";
        }
        if(res.size() && res.back()==' ')res.pop_back();
        return res;
    }
};