[Python] 486ms, beat 100%


  • 0

    Using the same technique: trie + DFS

    class Trie:
        def __init__(self, words):
            self.root = dict()
            for word in words:
                self.add(word)
    
        def add(self, word):
            node = self.root
            for char in word:
                if char not in node:
                    node[char] = dict()
                if '#' not in node:
                    node['#'] = []
                node['#'].append(word)
                node = node[char]
    
    class Solution(object):
        def wordSquares(self, words):
            """
            :type words: List[str]
            :rtype: List[List[str]]
            """
            trie = Trie(words)
            squares = []
            square = []
    
            def dfs(idx):
                if square and idx == len(square[0]):
                    squares.append(square[:])
                    return
                node = trie.root
                for word in square:
                    char = word[idx]
                    if char not in node:
                        return
                    node = node[char]
                for word in node['#']:
                    square.append(word)
                    dfs(idx + 1)
                    square.pop()
    
            dfs(0)
            return squares
    

Log in to reply
 

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