[lintcode] problem 634 – 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 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.

1
2
3
4
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

No.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).

No.2

Input:
[“abat”,”baba”,”atan”,”atal”]

Output:
[[“baba”,”abat”,”baba”,”atan”],[“baba”,”abat”,”baba”,”atal”]]

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
public class  {
private List<String> words = new ArrayList<>();
private Map<Character, TrieNode> children = new HashMap<>();
}

public class TrieTree {
private TrieNode root;

public TrieTree(TrieNode root) {
this.root = root;
}

public void insert(String word) {
TrieNode node = root;

for (char ch : word.toCharArray()) {
node.children.putIfAbsent(ch, new TrieNode());
node.children.get(ch).words.add(word);
node = node.children.get(ch);
}
}

public List<String> findByPrefix(String prefix) {
List<String> result = new ArrayList<>();
TrieNode node = root;

for (char ch : prefix.toCharArray()) {
if (!node.children.containsKey(ch))
return result;

node = node.children.get(ch);
}

result.addAll(node.words);
return result;
}
}

public List<List<String>> wordSquares(String[] words) {
List<List<String>> results = new ArrayList<>();

if (words == null || words.length == 0)
return results;

List<String> result = new ArrayList<>();
TrieTree trie = new TrieTree(new TrieNode());
int n = words[0].length();

for (String word : words)
trie.insert(word);

for (String word : words) {
result.add(word);
helper(results, result, trie, n);
result.remove(result.size() - 1);
}

return results;
}

private void helper(List<List<String>> results, List<String> result, TrieTree trie, int n) {
if (result.size() == n) {
results.add(new ArrayList(result));
return;
}

StringBuilder prefix = new StringBuilder();
int idx = result.size();

for (String res : result)
prefix.append(res.charAt(idx));

List<String> words = trie.findByPrefix(prefix.toString());

for (String word : words) {
result.add(word);
helper(results, result, trie, n);
result.remove(result.size() - 1);
}
}