45ms C++ solution. The code is a bit long though...


  • 0
    J
    struct TrieNode {
        TrieNode() {
            memset(child, 0, 26*(sizeof(TrieNode*)));
        }
        TrieNode* child[26];
    };
    
    class Solution {
    public:
        vector<vector<string>> wordSquares(vector<string>& words) {
            vector<vector<string> > ret;
            if(words.size() == 0)
            {
                return ret;
            }
    
            vector<TrieNode*> vt(words[0].length(), NULL);
            TrieNode* root = new TrieNode();
            for(int i=0;i<words.size();i++)
            {
                addWord(root, words[i]);
            }
            
            for(int i=0;i<words.size();i++)
            {
                bool skip = false;
                vector<string> vs;
                vs.push_back(words[i]);
                for(int j=1;j<words[i].length();j++)
                {
                    if(root->child[words[i][j] - 'a'])
                    {
                        vt[j] = root->child[words[i][j] - 'a'];
                        vs.push_back(string(1, words[i][j]));
                    }
                    else
                    {
                        skip = true;
                        break;
                    }
                }
                
                if(!skip)
                {
                    recur(ret, vt, vs, 1, 1);
                }
            }
            
            return ret;
        }
    private:
        void recur(vector<vector<string> > & ret, vector<TrieNode*> & vt, vector<string> & vs, int row, int col)
        {
            if(row == vt.size())
            {
                ret.push_back(vs);
                return;
            }
    
            for(int i=0;i<26;i++)
            {
                if(vt[row]->child[i] && vt[col]->child[i])
                {
                    TrieNode * r = vt[row];
                    TrieNode * c = vt[col];
                    vt[row] = r->child[i];
                    vs[row].push_back(i + 'a');
                    if(row != col)
                    {
                        vt[col] = c->child[i];
                        vs[col].push_back(i + 'a');
                    }
                    
                    if(col + 1 == vt.size())
                    {
                        recur(ret, vt, vs, row+1, row+1);
                    }
                    else
                    {
                        recur(ret, vt, vs, row, col+1);
                    }
                    
                    vt[row] = r;
                    vs[row].pop_back();
                    if(row != col)
                    {
                        vt[col] = c;
                        vs[col].pop_back();
                    }
                }
            }
            return;
        }
    
        void addWord(TrieNode * root, string s)
        {
            TrieNode* tn = root;
            for(int i=0;i<s.length();i++)
            {
                if(tn->child[s[i]-'a'] == NULL)
                {
                    tn->child[s[i]-'a'] = new TrieNode();
                }
                tn = tn->child[s[i]-'a'];
            }
            return;
        }
    };
    

Log in to reply
 

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