LeetCode 425. Word Squares

Description

https://leetcode.com/problems/word-squares/

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 k^th 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:

1. There are at least 1 and at most 1000 words.
2. All words will have the exact same length.
3. Word length is at least 1 and at most 5.
4. 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).

Explanation

use a trie to help find if prefix meets criteria

Python Solution

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_word = False
        self.word_list = []
        
class Trie:
    def __init__(self):
        self.root = TrieNode()
    
    def add(self, word):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
                
            node = node.children[char]
            node.word_list.append(word)
        node.is_word = True
    
    def find(self, word):
        node = self.root
        for char in word:
            node = node.children.get(char)
            if node is None:
                return None
            
        return node
    
    def word_prefix(self, prefix):
        node = self.find(prefix)
        return [] if node is None else node.word_list
            
            

class Solution:
    def wordSquares(self, words: List[str]) -> List[List[str]]:        
        
        trie = Trie()
        
        for word in words:
            trie.add(word)
        
        results = []
        
        for word in words:
            self.search(trie, [word], results)
            
        return results
    
    def search(self, trie, square, results):
        
        n, pos = len(square[0]), len(square)
        
        if pos == n:
            results.append(list(square))
            return
        
        for col in range(pos, n):
            prefix = ''.join(square[i][col] for i in range(pos))
            if trie.find(prefix) is None:
                return
            
        prefix = ''.join(square[i][pos] for i in range(pos))
        for word in trie.word_prefix(prefix):
            square.append(word)
            self.search(trie, square, results)
            square.pop()
            

• Time Complexity: ~N*26^L*L
• Space Complexity: ~NL + L * 1/2 = NL

where N is the number of input words and L is the length of a single word