struct TrieNode {
  vector<TrieNode*> children;
  vector<const string*> startsWith;
  TrieNode() : children(26) {}
  ~TrieNode() {
    for (TrieNode* child : children)
      delete child;
  }
};

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

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

 private:
  TrieNode root;

  void insert(const string& word) {
    TrieNode* node = &root;
    for (const char c : word) {
      const int i = c - 'a';
      if (!node->children[i])
        node->children[i] = new 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();
  }
}