Python Trie solution with comments

  • 1
    class Solution(object):
        def wordSquares(self, words):
            :type words: List[str]
            :rtype: List[List[str]]
            # collect all words in this subtrie
            def gather_trie(trie, res, item):
                if (not trie):
                    if (len(item) > 0):
                for c in trie:
                    gather_trie(trie[c], res, item + c)
            # add word to trie                
            def add_word(trie, word, start):
                if (start == len(word)):
                if (word[start] not in trie):
                    trie[word[start]] = {}
                add_word(trie[word[start]], word, start + 1)
            # build trie; dict of dicts
            # don't worry about leaf markers because all words are of same length 
            # and no duplicates
            def build_trie(words):
                trie = {}
                for word in words:
                    add_word(trie, word, 0)
                return trie
            # gather all possible words with given prefix
            def get_words(trie, prefix):
                for c in prefix:
                    if (c not in trie):
                        return []   # return empty if prefix not found
                    trie = trie[c]
                # traverse this subtrie to gather words
                res = []
                gather_trie(trie, res, prefix)
                return res
            # search through backtracking
            def backtrack(trie, res, item, chari, l, prefix):
                cands = get_words(trie, prefix)
                for cand in cands:
                    if (len(item) == l - 1):
                        res.append(list(item + [cand]))
                    # construct new prefix for next word search
                    newprefix = ''.join([word[chari + 1] for word in (item + [cand])])
                    backtrack(trie, res, item + [cand], chari + 1, l, newprefix)
            l = len(words)
            if (l == 0):
                return [[]]
            trie = build_trie(words)
            res = []
            backtrack(trie, res, [], 0, len(words[0]), '')
            return res

Log in to reply

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