Count Prefix and Suffix Pairs I
Easy · Tries · Explanation and code example
Trie
How it works
We can use a trie where each node is keyed by a pair of characters: one from the prefix and one from the suffix. By processing words in reverse order and storing them in this combined trie, when we look up a word, we find how many previously seen words have it as both prefix and suffix.
Steps
- Create a trie where each edge is labeled by a pair (prefix char, suffix char).
- Process words from the end of the array to the beginning.
- For each word, traverse the trie using pairs of (word[i], word[n-1-i]) and count matches.
- Then insert the word into the trie, incrementing counts at each node.
- Return the total count of matching pairs.
Complexity
- Time complexity:
O(n * m) - Space complexity:
O(n * m)
Where
nis the size of the input arraywords, andmis the maximum length of a string.
Things to watch for
Checking Only Prefix or Only Suffix
The condition requires words[i] to be BOTH a prefix AND a suffix of words[j]. Checking only one condition will give incorrect results.
Forgetting the Length Constraint
A string cannot be a prefix or suffix of a shorter string. Failing to check len(s1) <= len(s2) before comparing can lead to index out of bounds errors or false positives.
Python code example
from typing import List
class TrieNode:
def __init__(self):
self.children = {}
self.count = 0
class Trie:
def __init__(self):
self.root = TrieNode()
def add(self, w: str) -> None:
cur = self.root
for c1, c2 in zip(w, reversed(w)):
if (c1, c2) not in cur.children:
cur.children[(c1, c2)] = TrieNode()
cur = cur.children[(c1, c2)]
cur.count += 1
def count(self, w: str) -> int:
cur = self.root
for c1, c2 in zip(w, reversed(w)):
if (c1, c2) not in cur.children:
return 0
cur = cur.children[(c1, c2)]
return cur.count
class Solution:
def countPrefixSuffixPairs(self, words: List[str]) -> int:
res = 0
root = Trie()
for w in reversed(words):
res += root.count(w)
root.add(w)
return res
C++ code example
class TrieNode {
public:
unordered_map<string, TrieNode*> children;
int count = 0;
};
class Trie {
public:
TrieNode* root;
Trie() {
root = new TrieNode();
}
void add(string w) {
TrieNode* cur = root;
int n = w.size();
for (int i = 0; i < n; i++) {
string key = string(1, w[i]) + w[n - 1 - i];
if (cur->children.find(key) == cur->children.end()) {
cur->children[key] = new TrieNode();
}
cur = cur->children[key];
cur->count++;
}
}
int count(string w) {
TrieNode* cur = root;
int n = w.size();
for (int i = 0; i < n; i++) {
string key = string(1, w[i]) + w[n - 1 - i];
if (cur->children.find(key) == cur->children.end()) return 0;
cur = cur->children[key];
}
return cur->count;
}
};
class Solution {
public:
int countPrefixSuffixPairs(vector<string>& words) {
int res = 0;
Trie root;
for (int i = words.size() - 1; i >= 0; i--) {
res += root.count(words[i]);
root.add(words[i]);
}
return res;
}
};
Source and license
Adapted from NeetCode’s source · MIT license. Node types are supplied by LeetCode where needed.