Leetcode 425 Solution

class TrieNode
{
public:
    vector<TrieNode*> children;
    vector<int> indexes;
    TrieNode() : children(26, nullptr) {}
};

class Solution {
public:
    TrieNode* build(vector<string>& words)
    {
        TrieNode* root = new TrieNode();

        for (int i = 0; i < words.size(); i++)
        {
            auto word = words[i];
            auto p = root;

            for (auto ch : word)
            {
                if (p->children[ch - 'a'] == nullptr)
                    p->children[ch - 'a'] = new TrieNode();

                p = p->children[ch - 'a'];
                p->indexes.push_back(i);
            }
        }

        return root;
    }

    void search(vector<string>& words, vector<vector<string>>& res, vector<string>& curr, int level, TrieNode* root)
    {
        if (level >= curr[0].size())
        {
            res.push_back(curr);
            return;
        }

        string pre;
        for (int i = 0; i < level; i++)
            pre += curr[i][level];

        auto p = root;
        for (auto ch : pre)
        {
            if (p->children[ch - 'a'] == nullptr)
                return;
            p = p->children[ch - 'a'];
        }

        for (auto index : p->indexes)
        {
            curr.push_back(words[index]);
            search(words, res, curr, level + 1, root);
            curr.pop_back();
        }
    }

    vector<vector<string>> wordSquares(vector<string>& words) {
        TrieNode* root = build(words);

        vector<vector<string>> res;

        for (auto word : words)
        {
            vector<string> square;
            square.push_back(word);
            search(words, res, square, 1, root);
        }

        return res;
    }
};