题目地址:
https://www.lintcode.com/problem/word-squares/description
对于一个若干单词组成的数组,对于任意的 i i i,如果第 i i i行和第 i i i列念出来的单词是一样的,则称其为一个单词方阵。给定若干单词组成的数组 A A A,求所有 A A A中的单词组成的单词方阵。题目保证单词只含英文小写字母,长度都相同,并且没有重复单词。
思路是DFS + 剪枝。其实就是直接枚举每个位置可以填哪些单词即可,没有别的好的办法。但我们仍然可以做若干剪枝操作。首先,
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class Solution {
// 开一个前缀树的类,可以实现插入单词,和得到某个字符串为前缀的所有单词,这两个功能
class Trie {
class Node {
boolean isWord;
Node[] nexts;
public Node() {
nexts = new Node[26];
}
}
private Node root;
// key是字符串,value是以key为前缀的所有单词组成的list。该map是做缓存之用
private Map<String, List<String>> map;
public Trie() {
root = new Node();
map = new HashMap<>();
}
public Trie(String[] words) {
this();
for (String word : words) {
insert(word);
}
}
public void insert(String s) {
Node cur = root;
for (int i = 0; i < s.length(); i++) {
char ch = s.charAt(i);
if (cur.nexts[ch - 'a'] == null) {
cur.nexts[ch - 'a'] = new Node();
}
cur = cur.nexts[ch - 'a'];
}
cur.isWord = true;
}
// 得到以pref为前缀的所有单词
public List<String> getPrefix(String pref) {
// 如果缓存里有该前缀,则直接拿出来
if (map.containsKey(pref)) {
return map.get(pref);
}
List<String> res = new ArrayList<>();
StringBuilder sb = new StringBuilder();
Node cur = root;
for (int i = 0; i < pref.length(); i++) {
char ch = pref.charAt(i);
// 如果不存在该前缀的单词,则直接返回空列表
if (cur.nexts[ch - 'a'] == null) {
return res;
}
sb.append(ch);
cur = cur.nexts[ch - 'a'];
}
dfs(cur, sb, res);
map.putIfAbsent(pref, res);
return res;
}
private void dfs(Node cur, StringBuilder sb, List<String> res) {
if (cur.isWord) {
res.add(sb.toString());
}
for (char ch = 'a'; ch <= 'z'; ch++) {
if (cur.nexts[ch - 'a'] != null) {
sb.append(ch);
dfs(cur.nexts[ch - 'a'], sb, res);
sb.setLength(sb.length() - 1);
}
}
}
}
/*
* @param words: a set of words without duplicates
* @return: all word squares
*/
public List<List<String>> wordSquares(String[] words) {
// write your code here
if (words == null || words.length == 0) {
return res;
}
Trie trie = new Trie(words);
List<List<String>> res = new ArrayList<>();
dfs(0, new ArrayList<>(), words, res, trie);
return res;
}
// pos记录枚举到的列的下标,cur存已经枚举的单词
private void dfs(int pos, List<String> cur, String[] words, List<List<String>> res, Trie trie) {
// 如果所有列都已经枚举完了,说明找到了一个合法解,将其加入res
if (pos == words[0].length()) {
res.add(new ArrayList<>(cur));
return;
}
// 得到下一个要枚举的单词的前缀应该是什么
StringBuilder sb = new StringBuilder();
for (int i = 0; i < pos; i++) {
sb.append(cur.get(i).charAt(pos));
}
// 得到以该前缀开头的所有单词
List<String> nextWord = trie.getPrefix(sb.toString());
for (int i = 0; i < nextWord.size(); i++) {
// 判断一下,如果枚举了nextWord[i]这个单词的话,后面单词应该拥有的前缀是否存在对应的单词,
// 如果不存在,那直接不要枚举当前单词
if (!checkPref(pos, nextWord.get(i), cur, words[0].length(), trie)) {
continue;
}
// 枚举当前单词作为下一个单词
cur.add(nextWord.get(i));
dfs(pos + 1, cur, words, res, trie);
cur.remove(cur.size() - 1);
}
}
private boolean checkPref(int pos, String next, List<String> cur, int wordLen, Trie trie) {
// 算一下pos之后的列的竖着看的前缀
for (int i = pos + 1; i < wordLen; i++) {
StringBuilder sb = new StringBuilder();
for (int j = 0; j < pos; j++) {
sb.append(cur.get(j).charAt(i));
}
sb.append(next.charAt(i));
// 如果没有单词以sb为前缀,那返回false
if (trie.getPrefix(sb.toString()).isEmpty()) {
return false;
}
}
return true;
}
}
时空复杂度依赖于具体有哪些单词。