Word Squares

Introduction

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.

Trie Tree

Trie Tree is a data structure used to quickly find all elements with same prefix in O(log(N)) time. Each node in Trie tree represents a prefix and all of its children are also Trie node and represents a prefix that starts with prefix in parent node and extends with one more character. It’s faster to search for strings with a specific prefix than check the prefix of each string in a brute force method.

Reference source

Solution

The basic idea is to try every string in the list as the first string in word square, then prefix of the second string in word square is determined by the second character in first string. We search it in our Trie tree to find all strings with this prefix and apply each one as the second string in word square. When two strings are placed in word square, then the prefix of third string is determined by the third character of first and second string, then search Trie tree for all strings that meet the requirement. Repeatedly search for strings with new computed prefix until the word square is filled up. For this kind of backtracking problem, a recursive method that invoke itself is much easier to implement and more understandable than using stack or queue to search for an answer.

It perfectly passed all tests with clear logic, but doesn’t have a good running time, because TrieNode tree does not cache a list of strings for specific prefix, so they will be calculated every time, but makes memory usage at minimum and will not cause memory explosion if input string list is very large. Also, there are many deep copies of string list in search() function, which may further slow down the program.

Reference source

public class WordSquares {

  class TrieNode {
    public String prefix;
    public Map<String, TrieNode> children;
    public boolean isWord;

    public TrieNode() {
      prefix = "";
      children = new HashMap<>();
      isWord = false;
    }
  }

  // Find all strings with specific prefix by traversing TrieNode tree.
  public List<String> findByPrefix(TrieNode root, String prefix) {
    List<String> list = new ArrayList<>();
    Queue<TrieNode> q = new LinkedList<>();
    q.add(root);
    while(!q.isEmpty()) {
      TrieNode node = q.poll();
      if(node.prefix.startsWith(prefix)) {
        if(node.isWord) list.add(node.prefix);
        q.addAll(node.children.values());
      }
      else if(prefix.startsWith(node.prefix)) {
        q.addAll(node.children.values());
      }
    }
    return list;
  }

  // Create TrieNode tree to store all prefixes.
  public TrieNode init(String[] words) {
    TrieNode root = new TrieNode();
    for(String word : words) {
      TrieNode node = root;
      for(int i=0; i<word.length(); i++) {
        String prefix = word.substring(0, i+1);
        if(node.children.containsKey(prefix)) {
          node = node.children.get(prefix);
        }
        else {
          TrieNode newnode = new TrieNode();
          newnode.prefix = prefix;
          node.children.put(prefix, newnode);
          node = newnode;
        }
      }
      node.isWord = true;
    }
    return root;
  }

  // Recursively search until a result is reached.
  public void search(List<String> words, TrieNode root, List<List<String>> result) {
    int len = words.size();
    String prefix = "";
    for(String word : words) {
      prefix += word.charAt(len);
    }

    List<String> options = findByPrefix(root, prefix);
    for(String option : options) {
      List<String> list = new ArrayList<>(words);
      list.add(option);
      if(list.size() == option.length()) {
        result.add(list);
      }
      else {
        search(list, root, result);
      }
    }
  }

  public List<List<String>> wordSquares(String[] words) {
    List<List<String>> result = new ArrayList<>();
    // Deal with corner case when only one word in list.
    if(words.length == 1) {
      List<String> list = new ArrayList<>();
      list.add(words[0]);
      result.add(list);
      return result;
    }

    TrieNode root = init(words);   

    for(String word : words) {
    	List<String> list = new ArrayList<>();
    	list.add(word);
    	search(list, root, result);
    }
    return result;
  }

  public static void main(String[] args) {
    String[] words = new String[]{"area","lead","wall","lady","ball"};
    WordSquares ws = new WordSquares();
    List<List<String>> list = ws.wordSquares(words);
    for(List<String> sublist : list) {
      for(String str : sublist) {
        System.out.print(str + " ");
      }
      System.out.println();
    }
  }

}