Count Prefix and Suffix Pairs II
HardPrompt
You are given a 0-indexed string array words.
Let's define a boolean function isPrefixAndSuffix that takes two strings, str1 and str2:
isPrefixAndSuffix(str1, str2)returnstrueifstr1is both a prefix and a suffix ofstr2, andfalseotherwise.
For example, isPrefixAndSuffix("aba", "ababa") is true because "aba" is a prefix of "ababa" and also a suffix, but isPrefixAndSuffix("abc", "abcd") is false.
Return an integer denoting the number of index pairs (i, j) such that i < j, and isPrefixAndSuffix(words[i], words[j]) is true.
Example 1:
Input: words = ["a","aba","ababa","aa"]
Output: 4
Explanation: In this example, the counted index pairs are:
i = 0 and j = 1 because isPrefixAndSuffix("a", "aba") is true.
i = 0 and j = 2 because isPrefixAndSuffix("a", "ababa") is true.
i = 0 and j = 3 because isPrefixAndSuffix("a", "aa") is true.
i = 1 and j = 2 because isPrefixAndSuffix("aba", "ababa") is true.
Therefore, the answer is 4.Example 2:
Input: words = ["pa","papa","ma","mama"]
Output: 2
Explanation: In this example, the counted index pairs are:
i = 0 and j = 1 because isPrefixAndSuffix("pa", "papa") is true.
i = 2 and j = 3 because isPrefixAndSuffix("ma", "mama") is true.
Therefore, the answer is 2. Example 3:
Input: words = ["abab","ab"]
Output: 0
Explanation: In this example, the only valid index pair is i = 0 and j = 1, and isPrefixAndSuffix("abab", "ab") is false.
Therefore, the answer is 0.
Constraints:
1 <= words.length <= 1051 <= words[i].length <= 105words[i]consists only of lowercase English letters.- The sum of the lengths of all
words[i]does not exceed5 * 105.
Approaches
3 approaches with complexity analysis and trade-offs.
The most straightforward approach is to simulate the process described in the problem directly. We can use nested loops to generate every possible pair of indices (i, j) such that i < j. For each pair, we then check if words[i] is both a prefix and a suffix of words[j]. If it is, we increment a counter.
Algorithm
- Initialize a counter
countto 0. - Iterate through the
wordsarray with an outer loop for indexifrom 0 ton-1(wherenis the number of words). - Start an inner loop for index
jfromi+1ton-1. - Inside the inner loop, check if
words[i]is both a prefix and a suffix ofwords[j].- This can be done using built-in string functions like
startsWith()andendsWith(). - A necessary condition is that the length of
words[i]must not be greater than the length ofwords[j].
- This can be done using built-in string functions like
- If the condition holds true, increment the
count. - After the loops complete, return the total
count.
Walkthrough
This method involves a double loop. The outer loop iterates from the first word to the second-to-last word, picking words[i]. The inner loop iterates from words[i+1] to the last word, picking words[j]. For each pair (words[i], words[j]), we perform two checks: words[j].startsWith(words[i]) and words[j].endsWith(words[i]). If both checks return true, we've found a valid pair and we increment our result counter.
class Solution { public long countPrefixSuffixPairs(String[] words) { long count = 0; int n = words.length; for (int i = 0; i < n; i++) { for (int j = i + 1; j < n; j++) { // The length check is implicitly handled by startsWith/endsWith // but it's good practice to be aware of it. if (words[j].startsWith(words[i]) && words[j].endsWith(words[i])) { count++; } } } return count; }}Complexity
Time
O(N² * L), where N is the number of words and L is the maximum length of a word. For each of the O(N²) pairs, the `startsWith` and `endsWith` operations can take up to O(L) time.
Space
O(1) extra space, as we only use a few variables to keep track of loops and the count.
Trade-offs
Pros
Simple to understand and implement.
Requires no complex data structures.
Cons
Extremely inefficient for large inputs due to the nested loops.
The time complexity of O(N² * L) will cause a 'Time Limit Exceeded' error on platforms like LeetCode for the given constraints.
To improve upon the brute-force approach, we can change our perspective. Instead of picking a pair (i, j) and checking, let's process the words one by one. When we are at words[j], we need to efficiently count how many preceding words words[i] (with i < j) are a prefix-suffix of words[j]. We can use a hash map to store the counts of all words seen so far. For each new word, we find all its prefixes that are also suffixes and sum up their counts from the hash map.
Algorithm
- Initialize a counter
countto 0 and aHashMap<String, Integer>to store the frequency of words encountered so far. - Iterate through each
wordin thewordsarray. - For the current
word, iterate through all its possible prefixes, from length 1 up to its full length. - For each
prefix, check if thewordalso ends with thisprefix. - If it does, it means this
prefixis a valid prefix-suffix. Look up thisprefixin the frequency map and add its count to the totalcount. - After checking all prefixes for the current
word, update the frequency map with the currentworditself to make it available for subsequent words.
Walkthrough
We iterate through the words array. We maintain a hash map, freq, that stores each unique word we've seen and how many times it has appeared. For each word in the input array, we generate all of its prefixes. For each prefix, we check if it's also a suffix of the current word. If it is, we add the frequency of this prefix-string (which we get from our freq map) to our running total. After processing all prefixes of the current word, we add the word itself to the freq map (or increment its count) so it can be matched by later words in the array.
import java.util.HashMap;import java.util.Map; class Solution { public long countPrefixSuffixPairs(String[] words) { long count = 0; Map<String, Integer> freq = new HashMap<>(); for (String word : words) { int len = word.length(); // Check all prefixes of the current word for (int i = 1; i <= len; i++) { String prefix = word.substring(0, i); // Check if the prefix is also a suffix if (word.endsWith(prefix)) { count += freq.getOrDefault(prefix, 0); } } // Add the current word to the frequency map for future words to check against freq.put(word, freq.getOrDefault(word, 0) + 1); } return count; }}Complexity
Time
O(Σ Lᵢ²), where Lᵢ is the length of the i-th word. For each word of length L, we iterate through L prefixes, and for each, `substring` and `endsWith` take O(L) time. This results in an O(L²) process for each word.
Space
O(S), where S is the sum of the lengths of all unique words, to store the frequency map.
Trade-offs
Pros
More efficient than the pure brute-force approach.
Avoids the O(N²) iteration over pairs of indices.
Cons
The time complexity is still high, especially for words with great length, as checking all prefixes and their suffix condition takes O(L²) time for a word of length L.
This can lead to a 'Time Limit Exceeded' for inputs with very long strings.
The most efficient solution involves combining a Trie with a fast string-matching algorithm like KMP (or Z-algorithm). The core idea is to process words one by one, using the Trie to store words seen so far. For each new word, we query the Trie to find matching pairs. The KMP algorithm's preprocessing step (computing the LPS array) allows us to identify all of the word's prefixes that are also suffixes in linear time. This avoids the costly O(L²) check from the previous approach.
Algorithm
- Define a
TrieNodestructure containing an array/map for children and acountfield. - Initialize a Trie
rootand atotal_pairscounter. - Iterate through each
wordin thewordsarray:- Query Phase:
a. Compute the LPS (Longest Proper Prefix which is also Suffix) array for the
wordusing the KMP preprocessing algorithm. This takes O(length ofword). b. From the LPS array, determine all lengthskfor whichword[0...k-1]is a prefix-suffix ofword. Store these lengths in a hash set for O(1) lookup. c. Traverse the Trie with the currentword. For each charactercat indexi, move to the corresponding child node. d. If at any point the path doesn't exist in the Trie, stop the traversal. e. After moving to a new node corresponding to the prefix of lengthi+1, check ifi+1is one of the pre-calculated prefix-suffix lengths. If it is, add thecountstored at this Trie node tototal_pairs. - Update Phase:
a. Insert the current
wordinto the Trie. Traverse the path for theword, creating nodes if they don't exist. b. Increment thecountat the final node corresponding to the end of theword.
- Query Phase:
a. Compute the LPS (Longest Proper Prefix which is also Suffix) array for the
- Return
total_pairs.
Walkthrough
In this optimal approach, we iterate through the words array, and for each word, we first query for pairs and then update our data structure. The data structure is a Trie, where each node stores a count of how many words inserted so far end at that node.
Querying: For a given word, we need to find how many previously inserted words p are a prefix-suffix of word. We first find all prefixes of word that are also its suffixes. The KMP algorithm's LPS array is perfect for this. We can compute the LPS array for word in O(L) time. Then, by chaining back from lps[L-1], we can find the lengths of all such prefix-suffixes. With these lengths identified, we traverse the Trie with word. At each step i of the traversal (corresponding to prefix word[0...i-1]), if the length i+1 is a valid prefix-suffix length, we add the count from the current Trie node to our result. This single traversal effectively queries for all valid prefix-suffixes at once.
Updating: After the query, we insert the current word into the Trie by traversing its character path and incrementing the count of the final node.
This combination ensures that processing each word takes time proportional to its length, leading to an overall linear time complexity with respect to the total number of characters.
import java.util.HashSet;import java.util.Set; class TrieNode { TrieNode[] children; int count; TrieNode() { children = new TrieNode[26]; count = 0; }} class Solution { private TrieNode root; public long countPrefixSuffixPairs(String[] words) { root = new TrieNode(); long totalPairs = 0; for (String word : words) { totalPairs += query(word); insert(word); } return totalPairs; } private int[] computeLPS(String s) { int n = s.length(); int[] lps = new int[n]; int length = 0; int i = 1; while (i < n) { if (s.charAt(i) == s.charAt(length)) { length++; lps[i] = length; i++; } else { if (length != 0) { length = lps[length - 1]; } else { lps[i] = 0; i++; } } } return lps; } private long query(String word) { long count = 0; int[] lps = computeLPS(word); Set<Integer> psLengths = new HashSet<>(); int k = word.length(); while (k > 0) { psLengths.add(k); // For the next iteration, find the longest proper prefix-suffix of the current prefix-suffix k = lps[k - 1]; } TrieNode curr = root; for (int i = 0; i < word.length(); i++) { int charIndex = word.charAt(i) - 'a'; if (curr.children[charIndex] == null) { break; } curr = curr.children[charIndex]; if (psLengths.contains(i + 1)) { count += curr.count; } } return count; } private void insert(String word) { TrieNode curr = root; for (char c : word.toCharArray()) { int charIndex = c - 'a'; if (curr.children[charIndex] == null) { curr.children[charIndex] = new TrieNode(); } curr = curr.children[charIndex]; } curr.count++; }}Complexity
Time
O(S), where S is the sum of the lengths of all words. For each word of length L, computing the LPS array takes O(L), and querying/inserting into the Trie also takes O(L). Summing over all words gives a total time complexity linear in S.
Space
O(S), where S is the total length of all words. This space is used for storing the Trie. The LPS array and set for each word take O(L_max) space, where L_max is the maximum word length.
Trade-offs
Pros
Optimal time complexity, linear in the total number of characters.
Scales well to the maximum constraints of the problem.
Efficiently handles the prefix-and-suffix matching condition.
Cons
Significantly more complex to implement, requiring knowledge of both Tries and the KMP algorithm.
The constant factor might be higher than simpler solutions for very small inputs, but its scalability is far superior.
Solutions
Solution
class Node { Map < Integer , Node > children = new HashMap <>(); int cnt ; } class Solution { public long countPrefixSuffixPairs ( String [] words ) { long ans = 0 ; Node trie = new Node (); for ( String s : words ) { Node node = trie ; int m = s . length (); for ( int i = 0 ; i < m ; ++ i ) { int p = s . charAt ( i ) * 32 + s . charAt ( m - i - 1 ); node . children . putIfAbsent ( p , new Node ()); node = node . children . get ( p ); ans += node . cnt ; } ++ node . cnt ; } return ans ; } }class Node { public: unordered_map < int , Node *> children ; int cnt ; Node () : cnt ( 0 ) {} }; class Solution { public: long long countPrefixSuffixPairs ( vector < string >& words ) { long long ans = 0 ; Node * trie = new Node (); for ( const string & s : words ) { Node * node = trie ; int m = s . length (); for ( int i = 0 ; i < m ; ++ i ) { int p = s [ i ] * 32 + s [ m - i - 1 ]; if ( node -> children . find ( p ) == node -> children . end ()) { node -> children [ p ] = new Node (); } node = node -> children [ p ]; ans += node -> cnt ; } ++ node -> cnt ; } return ans ; } };class Node : __slots__ = [ "children" , "cnt" ] def __init__ ( self ): self . children = {} self . cnt = 0 class Solution : def countPrefixSuffixPairs ( self , words : List [ str ]) -> int : ans = 0 trie = Node () for s in words : node = trie for p in zip ( s , reversed ( s )): if p not in node . children : node . children [ p ] = Node () node = node . children [ p ] ans += node . cnt node . cnt += 1 return ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.