Skip to content

3045. Count Prefix and Suffix Pairs II 👍

  • Time: $O(\Sigma |\texttt{words[i]}|)$
  • Space: $O(\Sigma |\texttt{words[i]}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
struct TrieNode {
  unordered_map<int, shared_ptr<TrieNode>> children;
  int count = 0;
};

class Trie {
 public:
  int insert(const string& word) {
    const int n = word.length();
    int count = 0;
    shared_ptr<TrieNode> node = root;
    for (int i = 0; i < n; ++i) {
      const int j = hash(word[i], word[n - 1 - i]);
      if (node->children[j] == nullptr)
        node->children[j] = make_shared<TrieNode>();
      node = node->children[j];
      count += node->count;
    }
    ++node->count;
    return count;
  }

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

  static int hash(char prefix, char suffix) {
    return 26 * (prefix - 'a') + (suffix - 'a');
  }
};

class Solution {
 public:
  // Same as 3042. Count Prefix and Suffix Pairs I
  long long countPrefixSuffixPairs(vector<string>& words) {
    long ans = 0;
    Trie trie;

    for (const string& word : words)
      ans += trie.insert(word);

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class TrieNode {
  Map<Integer, TrieNode> children = new HashMap<>();
  int count = 0;
}

class Trie {
  public int insert(final String word) {
    final int n = word.length();
    int count = 0;
    TrieNode node = root;
    for (int i = 0; i < n; ++i) {
      final char prefix = word.charAt(i);
      final char suffix = word.charAt(n - 1 - i);
      final int key = (prefix - 'a') * 26 + (suffix - 'a');
      node.children.putIfAbsent(key, new TrieNode());
      node = node.children.get(key);
      count += node.count;
    }
    ++node.count;
    return count;
  }

  private TrieNode root = new TrieNode();
}

class Solution {
  // Same as 3045. Count Prefix and Suffix Pairs II
  public long countPrefixSuffixPairs(String[] words) {
    long ans = 0;
    Trie trie = new Trie();

    for (final String word : words)
      ans += trie.insert(word);

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class TrieNode:
  def __init__(self):
    self.children: dict[tuple[str, str], TrieNode] = {}
    self.count = 0


class Trie:
  def __init__(self):
    self.root = TrieNode()

  def insert(self, word: str) -> int:
    node = self.root
    count = 0
    for i, prefix in enumerate(word):
      suffix = word[-1 - i]
      node = node.children.setdefault((prefix, suffix), TrieNode())
      count += node.count
    node.count += 1
    return count


class Solution:
  # Same as 3045. Count Prefix and Suffix Pairs II
  def countPrefixSuffixPairs(self, words: list[str]) -> int:
    trie = Trie()
    return sum(trie.insert(word) for word in words)