C++ Trie solution w/o MLE


  • 0
    H

    As complained by many others, the Trie approach for C++ results in MLE for large test data. After some trial and error, I've figured out a way to make the Trie approach work w/o MLE. The trick is to use as much stack memory as possible before dynamically allocating Trie nodes on heap. Here is the code.

    class Trie {
    public:
        Trie() : m_word(false) {
            fill(m_children.begin(), m_children.end(), nullptr);
        }
            
        bool hasChild(char c) {
            return m_children[c - 'a'] != nullptr;
        }
        
        void addChild(char c, Trie* ptr) {
            m_children[c - 'a'] = ptr;
        }
        
        Trie* getChild(char c) {
            return m_children[c - 'a'];
        }
        
        bool isWord() {
            return m_word;
        }
        
        void setWord(bool _word) {
            m_word = _word;
        }
            
    private:
        array<Trie*, 26> m_children;
        bool m_word;
    };
    
    class Solution {
    public:
        vector<string> findAllConcatenatedWordsInADict(vector<string>& words) {
            array<Trie, 10000> trie_buffer1;
            vector<unique_ptr<Trie>> trie_buffer2;
            Trie root;
            int cnt = 0;
            for (const string &str : words) {
                Trie* cur = &root;
                for (const char c : str) {
                    if (!cur->hasChild(c)) {
                        if (cnt < 10000) {
                            cur->addChild(c, &trie_buffer1[cnt++]);
                        } else {
                            trie_buffer2.emplace_back(new Trie());
                            cur->addChild(c, trie_buffer2.back().get());
                        }
                    }
                    cur = cur->getChild(c);
                }
                cur->setWord(true);
            }
            
            vector<string> ret;
            for (const string &str : words) {
                const int n = str.size();
                vector<int> dp(n + 1, 0);
                dp[0] = 1;
                for (int i = 0; i <= n; ++i) {
                    if (!dp[i]) continue;
                    Trie* cur = &root;
                    for (int j = i + 1; j <= n; ++j) {
                        if (!cur->hasChild(str[j - 1])) break;
                        cur = cur->getChild(str[j - 1]);
                        if (cur->isWord()) dp[j] = 1;
                    }
                    if (i == 0) dp[n] = 0;
                    if (dp[n]) {
                        ret.push_back(str);
                        break;
                    }
                }
            }
            
            return ret;
        }    
    };
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.