Word Break II solution

Word-Break-2-se7so.cpp

struct TrieNode {
public:
    void insert(const string& word, int i) {
        
        if (i >= word.size()) {
            return;
        }
        
        if (children[word[i]] == NULL) {
            children[word[i]] = new TrieNode();
        }
            
        children[word[i]]->insert(word, i+1);
        
        if (i == word.size()-1) {
            children[word[i]]->endOfWord = true;
        }
    }
    
    TrieNode* search(const string& word, int i) {
        if (i == word.size() || children[word[i]] == NULL) {
            return NULL;
        }
        
        if (i == word.size()-1) {
            return children[word[i]];
        }
            
        return children[word[i]]->search(word, i+1);
    }
    
    bool endOfWord;
    TrieNode* children[512];
};

class Solution {
public:
    vector<string> wordBreak(string s, vector<string>& wordDict) {
        root = new TrieNode();
        
        for (string word : wordDict) {
            root->insert(word, 0);
        }
        
        return wordBreakHelper(s, 0);
    }
    
    vector<string> wordBreakHelper(const string& s, int i) {
        if (sol[i]) {
            return mem[i];
        }
        
        vector<string> res;
        
        if (i >= s.size()) {
            return res;
        }
        
        TrieNode* prev = root;
        string curWord = "";
        for (int j = i; j < s.size(); j++) {
            curWord += s[j];
            TrieNode* cur = prev->search(curWord, curWord.size()-1); // Search from last found char ;)
            
            if (cur != NULL && cur->endOfWord) {
                vector<string> perm = wordBreakHelper(s, j+1);
                
                if (perm.size() == 0 && j+1 >= s.size()) {
                    res.push_back(curWord);
                } else {
                    for (string w : perm) {
                        res.push_back(curWord + " " + w);
                    }
                }
            } else if (cur == NULL) {
                sol[i] = true;
                return (mem[i] = res);
            }
            
            prev = cur;
        }
        
        sol[i] = true;
        return (mem[i] = res);
    }
    
private:
    TrieNode* root;
    bool sol[5001];
    map<int, vector<string>> mem;
};