425. 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 row and column read the exact same string, where 0 ≤ < 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"
]
]
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"
]
]
The output consists of two word squares. The order of output does not matter (just the order of words in each word square matters).
Solution: Trie & DFS
Python:
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
Java:
public class Solution {
class TrieNode {
List<String> words;
TrieNode[] children;
TrieNode() {
words = new ArrayList<>();
children = new TrieNode[26];
}
}
class Trie {
TrieNode root;
Trie(String[] words) {
root = new TrieNode();
for (String word : words) {
add(word);
}
}
void add(String word) {
TrieNode node = root;
for (char ch : word.toCharArray()) {
int idx = ch - 'a';
if (node.children[idx] == null) {
node.children[idx] = new TrieNode();
}
node.words.add(word);
node = node.children[idx];
}
}
}
public List<List<String>> wordSquares(String[] words) {
Trie trie = new Trie(words);
List<List<String>> squares = new ArrayList<>();
List<String> square = new ArrayList<>();
dfs(trie, squares, square, 0);
return squares;
}
private void dfs(Trie trie, List<List<String>> squares, List<String> square, int idx) {
if (!square.isEmpty() && idx == square.get(0).length()) {
squares.add(new ArrayList<>(square));
return;
}
TrieNode node = trie.root;
for (String word : square) {
char ch = word.charAt(idx);
if (node.children[ch - 'a'] == null) return;
node = node.children[ch - 'a'];
}
for (String word : node.words) {
square.add(word);
dfs(trie, squares, square, idx + 1);
square.remove(square.size() - 1);
}
}
}