Welcome to Subscribe On Youtube

Formatted question description: https://leetcode.ca/all/425.html

425. Word Squares

Level

Hard

Description

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:

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

Solution

First construct a trie using all the words in words. Then find all word squares from the trie. A word square must have the same number of rows and columns and be symmetric. Use backtrack to find all word squares.

  • class Solution {
        public List<List<String>> wordSquares(String[] words) {
            TrieNode root = new TrieNode();
            for (String word : words)
                insert(root, word);
            List<List<String>> squares = new ArrayList<List<String>>();
            int length = words[0].length();
            String[] finished = new String[length];
            char[][] board = new char[length][length];
            search(root, root, board, 0, 0, length, finished, squares);
            return squares;
        }
    
        public void insert(TrieNode root, String word) {
            TrieNode node = root;
            int length = word.length();
            for (int i = 0; i < length; i++) {
                char letter = word.charAt(i);
                int index = letter - 'a';
                if (node.next[index] == null)
                    node.next[index] = new TrieNode();
                node = node.next[index];
            }
            node.word = word;
        }
        
        public void search(TrieNode current, TrieNode root, char[][] board, int row, int column, int length, String[] finished, List<List<String>> squares) {
            if (row == length)
                squares.add(new ArrayList<String>(Arrays.asList(finished)));
            else if (column == length) {
                finished[row] = current.word;
                search(root, root, board, row + 1, 0, length, finished, squares);
            } else if (column < row) {
                char letter = board[row][column];
                TrieNode nextNode = current.next[letter - 'a'];
                if (nextNode != null)
                    search(nextNode, root, board, row, column + 1, length, finished, squares);
            } else {
                for (char c = 'a'; c <= 'z'; c++) {
                    int index = c - 'a';
                    if (current.next[index] != null) {
                        if (row == 0 && root.next[index] == null)
                            continue;
                        board[row][column] = c;
                        board[column][row] = c;
                        search(current.next[index], root, board, row, column + 1, length, finished, squares);
                    }
                }
            }
        }
    }
    
    class TrieNode {
        TrieNode[] next;
        String word = null;
    
        public TrieNode() {
            next = new TrieNode[26];
        }
    }
    
  • // OJ: https://leetcode.com/problems/word-squares/
    // Time: O(MN)
    // Space: O(MN)
    struct TrieNode {
        TrieNode *next[26] = {};
        vector<int> indices;
    };
    class Solution {
        void addWord(TrieNode *node, string &s, int i) {
            node->indices.push_back(i);
            for (char c : s) {
                if (!node->next[c - 'a']) node->next[c - 'a'] = new TrieNode();
                node = node->next[c - 'a'];
                node->indices.push_back(i);
            }
        }
        vector<int> getWords(TrieNode *node, string &s) {
            for (char c : s) {
                node = node->next[c - 'a'];
                if (!node) return {};
            }
            return node->indices;
        }
    public:
        vector<vector<string>> wordSquares(vector<string>& A) {
            int N = A[0].size(), M = A.size();
            vector<vector<string>> ans;
            vector<string> tmp(N, string(N, ' '));
            TrieNode root;
            for (int i = 0; i < M; ++i) addWord(&root, A[i], i);
            function<void(int)> dfs = [&](int i) {
                if (i == N) {
                    ans.push_back(tmp);
                    return;
                }
                auto prefix = tmp[i].substr(0, i);
                for (auto &index : getWords(&root, prefix)) {
                    for (int k = i; k < N; ++k) {
                        tmp[i][k] = tmp[k][i] = A[index][k];
                    }
                    dfs(i + 1);
                }
            };
            dfs(0);
            return ans;
        }
    };
    
  • class Trie:
        def __init__(self):
            self.children = [None] * 26
            self.v = []
    
        def insert(self, w, i):
            node = self
            for c in w:
                idx = ord(c) - ord('a')
                if node.children[idx] is None:
                    node.children[idx] = Trie()
                node = node.children[idx]
                node.v.append(i)
    
        def search(self, w):
            node = self
            for c in w:
                idx = ord(c) - ord('a')
                if node.children[idx] is None:
                    return []
                node = node.children[idx]
            return node.v
    
    
    class Solution:
        def wordSquares(self, words: List[str]) -> List[List[str]]:
            def dfs(t):
                if len(t) == len(words[0]):
                    ans.append(t[:])
                    return
                idx = len(t)
                pref = [v[idx] for v in t]
                indexes = trie.search(''.join(pref))
                for i in indexes:
                    t.append(words[i])
                    dfs(t)
                    t.pop()
    
            trie = Trie()
            ans = []
            for i, w in enumerate(words):
                trie.insert(w, i)
            for w in words:
                dfs([w])
            return ans
    
    ############
    
    class Solution(object):
      def wordSquares(self, words):
        """
        :type words: List[str]
        :rtype: List[List[str]]
        """
    
        def dfs(path, res, m, prefix):
          if len(path) == m:
            res.append(path + [])
            return
    
          for word in prefix["".join(zip(*path)[len(path)])]:
            path.append(word)
            dfs(path, res, m, prefix)
            path.pop()
    
        if not words:
          return []
    
        prefix = collections.defaultdict(list)
        for word in words:
          for i in range(0, len(word)):
            prefix[word[:i]].append(word)
    
        m = len(words[0])
        res = []
        path = []
        for word in words:
          path.append(word)
          dfs(path, res, m, prefix)
          path.pop()
        return res
    
    
  • type Trie struct {
    	children [26]*Trie
    	v        []int
    }
    
    func newTrie() *Trie {
    	return &Trie{}
    }
    func (this *Trie) insert(word string, i int) {
    	node := this
    	for _, c := range word {
    		c -= 'a'
    		if node.children[c] == nil {
    			node.children[c] = newTrie()
    		}
    		node = node.children[c]
    		node.v = append(node.v, i)
    	}
    }
    func (this *Trie) search(word string) []int {
    	node := this
    	for _, c := range word {
    		c -= 'a'
    		if node.children[c] == nil {
    			return []int{}
    		}
    		node = node.children[c]
    	}
    	return node.v
    }
    
    func wordSquares(words []string) [][]string {
    	trie := newTrie()
    	for i, w := range words {
    		trie.insert(w, i)
    	}
    	ans := [][]string{}
    	var dfs func([]string)
    	dfs = func(t []string) {
    		if len(t) == len(words[0]) {
    			cp := make([]string, len(t))
    			copy(cp, t)
    			ans = append(ans, cp)
    			return
    		}
    		idx := len(t)
    		pref := []byte{}
    		for _, v := range t {
    			pref = append(pref, v[idx])
    		}
    		indexes := trie.search(string(pref))
    		for _, i := range indexes {
    			t = append(t, words[i])
    			dfs(t)
    			t = t[:len(t)-1]
    		}
    	}
    	for _, w := range words {
    		dfs([]string{w})
    	}
    	return ans
    }
    

All Problems

All Solutions