Short Python/C++ solution


  • 52

    Python Solution (accepted in ~870 ms)

    def wordSquares(self, words):
        n = len(words[0])
        fulls = collections.defaultdict(list)
        for word in words:
            for i in range(n):
                fulls[word[:i]].append(word)
        def build(square):
            if len(square) == n:
                squares.append(square)
                return
            for word in fulls[''.join(zip(*square)[len(square)])]:
                build(square + [word])
        squares = []
        for word in words:
            build([word])
        return squares
    

    Explanation

    I try every word for the first row. For each of them, try every fitting word for the second row. And so on. The first few rows determine the first few columns and thus determine how the next row's word must start. For example:

    wall      Try words      wall                     wall                      wall
    a...   => starting  =>   area      Try words      area                      area
    l...      with "a"       le..   => starting  =>   lead      Try words       lead
    l...                     la..      with "le"      lad.   => starting   =>   lady
                                                                with "lad"
    

    For quick lookup, my fulls dictionary maps prefixes to lists of words who have that prefix.


    C++ Solution (accepted in ~180 ms)

    class Solution {
    public:
        vector<vector<string>> wordSquares(vector<string>& words) {
            n = words[0].size();
            square.resize(n);
            for (string word : words)
                for (int i=0; i<n; i++)
                    fulls[word.substr(0, i)].push_back(word);
            build(0);
            return squares;
        
        }
        int n;
        unordered_map<string, vector<string>> fulls;
        vector<string> square;
        vector<vector<string>> squares;
        void build(int i) {
            if (i == n) {
                squares.push_back(square);
                return;
            }
            string prefix;
            for (int k=0; k<i; k++)
                prefix += square[k][i];
            for (string word : fulls[prefix]) {
                square[i] = word;
                build(i + 1);
            }
        }
    };
    

  • 0
    T

    Would fulls occupy too much space?
    Well, maybe the words are too short, so it is not necessary to use a trie structure.


  • 1

    Nice clear solution for an upvote. Maybe use unordered_map<string, vector<int>> to store word index instead of words themselves is better for space saving (?).

    Also, could for (string word : fulls[prefix]) cause a same word being used multiple times in a word square? Or I guess it is allowed to do so building word squares (even though the given word list contains no duplicates).

    Actually, I was thinking about this way to try every possible word as first role and then look for words with specific prefix. But somehow I though this is "brute force" and didn't attempt to implement. But the given constraint word length <= 5 and number of words <= 1000 should give a hint that it will be "brutal" to solve this problem anyway, isn't it? :-)

    In fact, for some problems in OJ, I had the "brute force" idea (e.g., max point on a line, Ones and Zeroes, etc) and tried to be clever but with no better result then gave up with nothing. I guess it is indeed critical to develop an instinct to have a rough idea about what the optimal complexity should be for a given problem. For me, how to implement a "brute force" algorithm could be challenging as well sometimes. For some problems, I know better if I am on the "right track" as long as I simply checked whether the popular solutions are long or short.


  • 0

    @TeXnician Just because TRIE implementation is a little bit long. Stefan always prefers shorter code.


  • 1

    loving defaultdict, very convenient factory initializor


  • 1

    @TeXnician I think TRIE may save some space here but finding prefix is slow using TRIE since it has to search the tree again and again. If you save all prefix when building the TRIE, then space complexity will be the same.


  • 0
    C

    Concise and excellent ! As always !


  • 0
    T

    @StefanPochmann Simply some unimaginable idea!


  • 0
    L

    This is unbelievably genius


  • 0
    W

    @StefanPochmann said in Short Python/C++ solution:
    how is for word in fulls[''.join(zip(*square)[len(square)])]: working?


Log in to reply
 

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