my python solution with trie and backtracking


  • 0
    I
    class TrieNode(object):
        def __init__(self):
            self.children = {}
            self.is_word = False
            
    class Trie(object):
        def __init__(self):
            self.root = TrieNode()
            
        def add_word(self, word):
            curr_node = self.root
            for char in word:
                if char not in curr_node.children:
                    curr_node.children[char] = TrieNode()
                curr_node = curr_node.children[char]
            curr_node.is_word = True
            
        def get_words_util(self, curr_node, curr_word, words):
            if curr_node.is_word:
                words.append(curr_word)
            else:
                for child_letter, child in curr_node.children.iteritems():
                    self.get_words_util(child, curr_word + child_letter, words)
            
        def get_words_with_prefix(self, prefix):
            #First find last node of prefix
            curr_node = self.root
            for char in prefix:
                if char not in curr_node.children:
                    return []
                    
                curr_node = curr_node.children[char]
                
            #now we have the last node, dfs to get all words
            words = []
            self.get_words_util(curr_node, prefix, words)
            return words
    
    class Solution(object):
        def gen_squares(self, words, curr_sq, trie, squares, rows_needed):
            if len(curr_sq) == rows_needed:
                squares.append(list(curr_sq))
            else:
                prefix = ''
                for row in xrange(len(curr_sq)):
                    prefix += curr_sq[row][len(curr_sq)]
                    
                words_with_prefix = trie.get_words_with_prefix(prefix)
                for word in words_with_prefix:
                    curr_sq.append(word)
                    self.gen_squares(words, curr_sq, trie, squares, rows_needed)
                    curr_sq.pop()
        
        def wordSquares(self, words):
            """
            :type words: List[str]
            :rtype: List[List[str]]
            """
            trie = Trie()
            word_len = len(words[0])
            for word in words:
                trie.add_word(word)
                
            squares = []
            for word in words:
                self.gen_squares(words, [word], trie, squares, word_len)
                
            return squares

Log in to reply
 

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