LeetCode: 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).

/**
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).
*/

// 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.
import java.util.*;
public class WordSquares {
    public static List<List<String>> wordSquares(String[] words) {
        List<List<String>> res = new ArrayList<List<String>>();
        if(words.length==0 || words[0].length()==0) {
            return res;
        }
        Map<String, Set<String>> map = new HashMap<>();
        int squareLen = words[0].length();
        // create all prefix from each word and put the word in hashset for the prefix key
        for(int i=0;i<words.length;i++){
for(int j=0;j<words[0].length();j++){
                if(!map.containsKey(words[i].substring(0, j))) {
                    map.put(words[i].substring(0, j), new HashSet<String>());
                }
                map.get(words[i].substring(0, j)).add(words[i]);
            }
        }
        helper(res, new ArrayList<String>(), 0, squareLen, map);
        return res;
    }

/**
the helper method will check the match of the prefix, if the match reaches the size of the square, we found one combination
*/
    public static void helper(List<List<String>> res, List<String> cur, int matched, int total, Map<String, Set<String>> map){
        if(matched == total) {
            res.add(new ArrayList<String>(cur));
            return;
        }
        // build search string
        StringBuilder sb = new StringBuilder();
        for(int i=0;i<matched;i++) {
            sb.append(cur.get(i).charAt(matched));
        }
//find all the words that match to this prefix string
        // backtracking
        Set<String> cand = map.get(sb.toString());
        if(cand==null) {
            return;
        }
        for(String str:cand){
            cur.add(str);
            helper(res, cur, matched+1, total, map);
            cur.remove(cur.size()-1);
        }
    }

public static void main(String [] args){
String[] words = {"area","lead","wall","lady","ball"};
List<List<String>> wordsList = wordSquares(words);
for(List<String> list: wordsList){
for(String s:list){
System.out.println(s);
}
System.out.println("");
}
words = new String[]{"abat","baba","atan","atal"};
wordsList = wordSquares(words);
for(List<String> list: wordsList){
for(String s:list){
System.out.println(s);
}
System.out.println("");
}
}
}