[LintCode] Problem 634 - Word Squares

Posted on | In ,

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 TrieNode {
    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);
    }
}