425. Word Squares

• Time:
• Space:

C++JavaPython

struct TrieNode {
  vector<shared_ptr<TrieNode>> children;
  vector<const string*> startsWith;
  TrieNode() : children(26) {}
};

class Trie {
 public:
  Trie(const vector<string>& words) {
    for (const string& word : words)
      insert(word);
  }

  vector<const string*> findBy(const string& prefix) {
    shared_ptr<TrieNode> node = root;
    for (const char c : prefix) {
      const int i = c - 'a';
      if (node->children[i] == nullptr)
        return {};
      node = node->children[i];
    }
    return node->startsWith;
  }

 private:
  shared_ptr<TrieNode> root = make_shared<TrieNode>();

  void insert(const string& word) {
    shared_ptr<TrieNode> node = root;
    for (const char c : word) {
      const int i = c - 'a';
      if (node->children[i] == nullptr)
        node->children[i] = make_shared<TrieNode>();
      node = node->children[i];
      node->startsWith.push_back(&word);
    }
  }
};

class Solution {
 public:
  vector<vector<string>> wordSquares(vector<string>& words) {
    if (words.empty())
      return {};

    const int n = words[0].length();
    vector<vector<string>> ans;
    vector<string> path;
    Trie trie(words);

    for (const string& word : words) {
      path.push_back(word);
      dfs(trie, n, path, ans);
      path.pop_back();
    }

    return ans;
  }

 private:
  void dfs(Trie& trie, const int n, vector<string>& path,
           vector<vector<string>>& ans) {
    if (path.size() == n) {
      ans.push_back(path);
      return;
    }

    const string prefix = getPrefix(path);

    for (const string* s : trie.findBy(prefix)) {
      path.push_back(*s);
      dfs(trie, n, path, ans);
      path.pop_back();
    }
  }

  // e.g. path = ["wall",
  //              "area"]
  //    prefix =  "le.."
  string getPrefix(const vector<string>& path) {
    string prefix;
    const int index = path.size();
    for (const string& s : path)
      prefix += s[index];
    return prefix;
  }
};

class TrieNode {
  public TrieNode[] children = new TrieNode[26];
  public List<String> startsWith = new ArrayList<>();
}

class Trie {
  public Trie(final String[] words) {
    for (final String word : words)
      insert(word);
  }

  public List<String> findBy(final String prefix) {
    TrieNode node = root;
    for (final char c : prefix.toCharArray()) {
      final int i = c - 'a';
      if (node.children[i] == null)
        return new ArrayList<>();
      node = node.children[i];
    }
    return node.startsWith;
  }

  private TrieNode root = new TrieNode();

  private void insert(final String word) {
    TrieNode node = root;
    for (final char c : word.toCharArray()) {
      final int i = c - 'a';
      if (node.children[i] == null)
        node.children[i] = new TrieNode();
      node = node.children[i];
      node.startsWith.add(word);
    }
  }
}

class Solution {
  public List<List<String>> wordSquares(String[] words) {
    if (words.length == 0)
      return new ArrayList<>();

    final int n = words[0].length();
    List<List<String>> ans = new ArrayList<>();
    List<String> path = new ArrayList<>();
    Trie trie = new Trie(words);

    for (final String word : words) {
      path.add(word);
      dfs(trie, n, path, ans);
      path.remove(path.size() - 1);
    }

    return ans;
  }

  private void dfs(Trie trie, final int n, List<String> path, List<List<String>> ans) {
    if (path.size() == n) {
      ans.add(new ArrayList<>(path));
      return;
    }

    final String prefix = getPrefix(path);

    for (final String s : trie.findBy(prefix)) {
      path.add(s);
      dfs(trie, n, path, ans);
      path.remove(path.size() - 1);
    }
  }

  // e.g. path = ["wall",
  //              "area"]
  //    prefix =  "le.."
  private String getPrefix(List<String> path) {
    StringBuilder sb = new StringBuilder();
    final int index = path.size();
    for (final String s : path)
      sb.append(s.charAt(index));
    return sb.toString();
  }
}

class TrieNode:
  def __init__(self):
    self.children: dict[str, TrieNode] = {}
    self.startsWith: list[str] = []


class Trie:
  def __init__(self, words: list[str]):
    self.root = TrieNode()
    for word in words:
      self._insert(word)

  def findBy(self, prefix: str) -> list[str]:
    node = self.root
    for c in prefix:
      if c not in node.children:
        return []
      node = node.children[c]
    return node.startsWith

  def _insert(self, word: str) -> None:
    node = self.root
    for c in word:
      node = node.children.setdefault(c, TrieNode())
      node.startsWith.append(word)


class Solution:
  def wordSquares(self, words: list[str]) -> list[list[str]]:
    if not words:
      return []

    n = len(words[0])
    ans = []
    path = []
    trie = Trie(words)

    for word in words:
      path.append(word)
      self._dfs(trie, n, path, ans)
      path.pop()

    return ans

  def _dfs(self, trie: Trie, n: int, path: list[str], ans: list[list[str]]):
    if len(path) == n:
      ans.append(path.copy())
      return

    prefix = self._getPrefix(path)

    for s in trie.findBy(prefix):
      path.append(s)
      self._dfs(trie, n, path, ans)
      path.pop()

  def _getPrefix(self, path: list[str]) -> str:
    """
    e.g. path = ["wall",
                 "area"]
       prefix =  "le.."
    """
    prefix = []
    index = len(path)
    for s in path:
      prefix.append(s[index])
    return ''.join(prefix)