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

All Problems

All Solutions