# coding=utf-8
# Given a set of words (without duplicates), find all word squares you can build from them.
#
# A sequence of words forms a valid word square if the kth row and column read the exact same string, where
# 0 ≤ k < max(numRows, numColumns).
#
# For example, the word sequence ["ball","area","lead","lady"] forms a word square because each word reads the same
# both horizontally and vertically.
#
# b a l l
# a r e a
# l e a d
# l a d y
# Note:
# There are at least 1 and at most 1000 words.
# All words will have the exact same length.
# Word length is at least 1 and at most 5.
# Each word contains only lowercase English alphabet a-z.
# Example 1:
#
# Input:
# ["area","lead","wall","lady","ball"]
#
# Output:
# [
#     [ "wall",
#       "area",
#       "lead",
#       "lady"
#       ],
#     [ "ball",
#       "area",
#       "lead",
#       "lady"
#       ]
# ]
#
# Explanation:
# The output consists of two word squares. The order of output does not matter (just the order of words in each word
#  square matters).
# Example 2:
#
# Input:
# ["abat","baba","atan","atal"]
#
# Output:
# [
#     [ "baba",
#       "abat",
#       "baba",
#       "atan"
#       ],
#     [ "baba",
#       "abat",
#       "baba",
#       "atal"
#       ]
# ]
#
# Explanation:
# The output consists of two word squares. The order of output does not matter (just the order of words in each word
# square matters).
import copy


class TrieNode(object):
    def __init__(self):
        self.children = {}
        self.start_with = set([])
        self.end_of_word = False


class Trie(object):
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        n = len(word)
        cur = self.root
        for i in range(n):
            if word[i] not in cur.children:
                cur.children[word[i]] = TrieNode()
            cur.start_with.add(word)
            cur = cur.children[word[i]]
        cur.end_of_word = True

    def find_by_prefix(self, prefix):
        cur = self.root
        for char in prefix:
            if char not in cur.children:
                return []
            else:
                cur = cur.children[char]
        return cur.start_with


class Solution(object):
    def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
        trie = Trie()
        for word in words:
            trie.insert(word)

        result = []
        interim_result = []
        for word in words:
            interim_result.append(word)
            self.search(len(word), trie, result, interim_result)
            interim_result.pop()
        return result

    def search(self, length, trie, result, interim_result):
        i = len(interim_result)
        if i == length:
            result.append(copy.copy(interim_result))
            return

        prefix = ''
        for word in interim_result:
            prefix += word[i]

        prefix_words = trie.find_by_prefix(prefix)
        for word in prefix_words:
            interim_result.append(word)
            self.search(length, trie, result, interim_result)
            interim_result.pop()

