Count Prefix and Suffix Pairs I

Difficulty: Easy | Acceptance: 77.80% | Paid: No Topics: Array, String, Trie, Rolling Hash, String Matching, Hash Function

You are given a 0-indexed array words of size n consisting of distinct strings.

The string words[i] can be paired with the string words[j] if:

• 0 <= i < j < n, and
• words[i] is equal to words[j]’s prefix and words[i] is equal to words[j]’s suffix.

Return the number of pairs (i, j) such that words[i] can be paired with words[j].

• Examples
• Constraints
• Brute Force
• Trie
• Rolling Hash

Examples

Example 1

Input: words = ["a","aba","ababa","aa"]
Output: 4
Explanation: In this example, the counted index pairs are:
(0,1) because words[0] = "a" is a prefix and suffix of words[1] = "aba"
(0,2) because words[0] = "a" is a prefix and suffix of words[2] = "ababa"
(1,2) because words[1] = "aba" is a prefix and suffix of words[2] = "ababa"
(3,3) is invalid because i < j is not satisfied.
Wait, (3,3) is not valid. Let's re-verify.
Actually, the pairs are (0,1), (0,2), (0,3), (1,2).
(0,3): "a" is prefix and suffix of "aa".
So there are 4 pairs.

Example 2

Input: words = ["pa","papa","ma","mama"]
Output: 2
Explanation: In this example, the counted index pairs are:
(0,1) because words[0] = "pa" is a prefix and suffix of words[1] = "papa"
(2,3) because words[2] = "ma" is a prefix and suffix of words[3] = "mama"

Example 3

Input: words = ["abab","ab"]
Output: 0
Explanation: In this example, the only index pair is (0,1), and words[0] is not a suffix of words[1].

Constraints

1 <= words.length <= 50
1 <= words[i].length <= 10
words[i] consists only of lowercase English letters.
words[i] are distinct.

Brute Force

Intuition Since the constraints are very small (words.length <= 50 and words[i].length <= 10), we can simply check every possible pair (i, j) where i < j. For each pair, we verify if words[i] is both a prefix and a suffix of words[j].

Steps

• Initialize a counter to 0.
• Iterate through the array with index i from 0 to n-1.
• Iterate through the array with index j from i+1 to n-1.
• For each pair (i, j), check if words[j] starts with words[i] AND ends with words[i].
• If both conditions are true, increment the counter.
• Return the counter.

class Solution:
    def countPrefixSuffixPairs(self, words: List[str]) -> int:
        count = 0
        n = len(words)
        for i in range(n):
            for j in range(i + 1, n):
                w1 = words[i]
                w2 = words[j]
                if w2.startswith(w1) and w2.endswith(w1):
                    count += 1
        return count

Complexity

• Time: O(n² * L), where n is the number of words and L is the maximum length of a word. We check n² pairs, and each check takes O(L) time.
• Space: O(1), we only use a few variables for counting.
• Notes: This is the most straightforward solution and is perfectly acceptable given the constraints.

Trie

Intuition We can optimize the lookup of previous words using a Trie. We iterate through the words. For each word words[j], we check all its prefixes. If a prefix is also a suffix of words[j], we check if that prefix string exists in our Trie (which contains all words[i] where i < j). If it does, we increment our count. Finally, we insert words[j] into the Trie for future checks.

Steps

• Initialize a Trie and a counter.
• Iterate through each word in the input array.
• For the current word w, iterate through all possible lengths k from 1 to len(w).
• Extract the prefix of length k and the suffix of length k.
• If the prefix equals the suffix, search for this string in the Trie. If found, increment the counter.
• Insert the current word w into the Trie.
• Return the total count.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False

class Solution:
    def countPrefixSuffixPairs(self, words: List[str]) -> int:
        root = TrieNode()
        count = 0
        
        def insert(word):
            node = root
            for ch in word:
                if ch not in node.children:
                    node.children[ch] = TrieNode()
                node = node.children[ch]
            node.is_end = True
        
        def search(word):
            node = root
            for ch in word:
                if ch not in node.children:
                    return False
                node = node.children[ch]
            return node.is_end

        for w in words:
            for k in range(1, len(w) + 1):
                prefix = w[:k]
                suffix = w[-k:]
                if prefix == suffix:
                    if search(prefix):
                        count += 1
            insert(w)
        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 word, we check L prefixes, and each check/search takes O(L).
• Space: O(n * L) to store the Trie.
• Notes: While more complex to implement, this approach scales better for larger inputs compared to Brute Force.

Rolling Hash

Intuition We can use a polynomial rolling hash to represent strings as integers. We store the hash values of all previously seen words in a set. For each new word, we calculate the hash of all its valid prefix-suffix candidates. If a hash exists in the set, we have a valid pair.

Steps

• Initialize a set to store hashes of seen words and a counter.
• Iterate through each word in the array.
• For the current word w, iterate through all possible lengths k from 1 to len(w).
• Check if the prefix of length k equals the suffix of length k.
• If they are equal, calculate the rolling hash of this substring.
• If the hash exists in the set, increment the counter.
• Calculate the hash of the full word w and add it to the set.
• Return the total count.

class Solution:
    def countPrefixSuffixPairs(self, words: List[str]) -> int:
        seen = set()
        count = 0
        base = 31
        mod = 10**9 + 7
        
        def get_hash(s):
            h = 0
            for ch in s:
                h = (h * base + ord(ch) - ord('a') + 1) % mod
            return h

        for w in words:
            for k in range(1, len(w) + 1):
                prefix = w[:k]
                suffix = w[-k:]
                if prefix == suffix:
                    if get_hash(prefix) in seen:
                        count += 1
            seen.add(get_hash(w))
        return count

Complexity

• Time: O(n * L²), where n is the number of words and L is the maximum length of a word. We iterate through substrings and compute hashes.
• Space: O(n) to store the hash values.
• Notes: This approach is efficient for string matching problems. Note that in production code, one must handle potential hash collisions (e.g., by double hashing or storing the string itself), but for this problem, a single hash is sufficient.