Count Prefix and Suffix Pairs 2

Problem

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) returns true if str1 is both a prefix and a suffix of str2, and false otherwise.

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.

Examples

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 <= 10^5
1 <= words[i].length <= 10^5
words[i] consists only of lowercase English letters.
The sum of the lengths of all words[i] does not exceed 5 * 105.

Solution

Method 1 – Trie with Prefix-Suffix Pairing

Intuition

We build a trie where each node represents a pair of characters: one from the prefix and one from the suffix (traversed in reverse). This allows us to efficiently count how many times a prefix-suffix pair has occurred so far. As we process each word, we traverse the trie, updating the count and accumulating the result whenever a matching prefix-suffix pair is found.

Approach

Initialize a trie root node and a result counter.
For each word in the input list:
- Reverse the word to get the suffix characters.
- For each character position i, use the pair (word[i], reversed_word[i]) as the key.
- Traverse or create the corresponding child node in the trie.
- Add the count at the current node to the result.
- After traversing the word, increment the count at the last node.
Return the accumulated result.

Example

Suppose words = ["ab", "ba", "aba", "bab"]

For "ab": pairs are (a, b), (b, a)
For "ba": pairs are (b, a), (a, b)
For "aba": (a, a), (b, b), (a, a)
For "bab": (b, b), (a, a), (b, b)

Complexity

⏰ Time complexity: O(n * m), where n is the number of words and m is the maximum word length.
🧺 Space complexity: O(n * m), for the trie structure.

Code

C++

#include <unordered_map>
#include <vector>
#include <string>
using namespace std;

struct TrieNode {
    unordered_map<pair<char, char>, TrieNode*, hash<pair<char, char>>> children;
    int count = 0;
};

class Solution {
public:
    int countPrefixSuffixPairs(vector<string>& words) {
        auto root = new TrieNode();
        int res = 0;
        for (const string& w : words) {
            TrieNode* node = root;
            string rev = w;
            reverse(rev.begin(), rev.end());
            for (int i = 0; i < w.size(); ++i) {
                pair<char, char> key = {w[i], rev[i]};
                if (!node->children.count(key)) node->children[key] = new TrieNode();
                node = node->children[key];
                res += node->count;
            }
            node->count++;
        }
        return res;
    }
};

Go

import "strings"

type TrieNode struct {
    children map[[2]byte]*TrieNode
    count    int
}

func countPrefixSuffixPairs(words []string) int {
    root := &TrieNode{children: make(map[[2]byte]*TrieNode)}
    res := 0
    for _, w := range words {
        node := root
        rev := reverse(w)
        for i := 0; i < len(w); i++ {
            key := [2]byte{w[i], rev[i]}
            if node.children[key] == nil {
                node.children[key] = &TrieNode{children: make(map[[2]byte]*TrieNode)}
            }
            node = node.children[key]
            res += node.count
        }
        node.count++
    }
    return res
}

func reverse(s string) string {
    runes := []rune(s)
    for i, j := 0, len(runes)-1; i < j; i, j = i+1, j-1 {
        runes[i], runes[j] = runes[j], runes[i]
    }
    return string(runes)
}

Java

import java.util.*;

class Solution {
    static class TrieNode {
        Map<String, TrieNode> children = new HashMap<>();
        int count = 0;
    }
    public int countPrefixSuffixPairs(String[] words) {
        TrieNode root = new TrieNode();
        int res = 0;
        for (String w : words) {
            TrieNode node = root;
            String rev = new StringBuilder(w).reverse().toString();
            for (int i = 0; i < w.length(); i++) {
                String key = w.charAt(i) + "," + rev.charAt(i);
                node.children.putIfAbsent(key, new TrieNode());
                node = node.children.get(key);
                res += node.count;
            }
            node.count++;
        }
        return res;
    }
}

Python

from collections import defaultdict
class TrieNode:
    def __init__(self):
        self.children = defaultdict(TrieNode)
        self.count = 0

class Solution:
    def countPrefixSuffixPairs(self, words: list[str]) -> int:
        root = TrieNode()
        res = 0
        for w in words:
            node = root
            rev = w[::-1]
            for i in range(len(w)):
                key = (w[i], rev[i])
                node = node.children[key]
                res += node.count
            node.count += 1
        return res

Kotlin

class Solution {
    class TrieNode {
        val children = mutableMapOf<Pair<Char, Char>, TrieNode>()
        var count = 0
    }
    fun countPrefixSuffixPairs(words: Array<String>): Int {
        val root = TrieNode()
        var res = 0
        for (w in words) {
            var node = root
            val rev = w.reversed()
            for (i in w.indices) {
                val key = Pair(w[i], rev[i])
                node = node.children.getOrPut(key) { TrieNode() }
                res += node.count
            }
            node.count++
        }
        return res
    }
}

TypeScript

class TrieNode {
    children: Map<string, TrieNode> = new Map();
    count = 0;
}

function countPrefixSuffixPairs(words: string[]): number {
    const root = new TrieNode();
    let res = 0;
    for (const w of words) {
        let node = root;
        const rev = w.split('').reverse().join('');
        for (let i = 0; i < w.length; i++) {
            const key = w[i] + ',' + rev[i];
            if (!node.children.has(key)) node.children.set(key, new TrieNode());
            node = node.children.get(key)!;
            res += node.count;
        }
        node.count++;
    }
    return res;
}

Rust

use std::collections::HashMap;

struct TrieNode {
    children: HashMap<(char, char), Box<TrieNode>>,
    count: usize,
}

impl TrieNode {
    fn new() -> Self {
        TrieNode { children: HashMap::new(), count: 0 }
    }
}

impl Solution {
    pub fn count_prefix_suffix_pairs(words: Vec<String>) -> i32 {
        let mut root = TrieNode::new();
        let mut res = 0;
        for w in &words {
            let mut node = &mut root;
            let rev: Vec<char> = w.chars().rev().collect();
            let chars: Vec<char> = w.chars().collect();
            for i in 0..w.len() {
                let key = (chars[i], rev[i]);
                node = node.children.entry(key).or_insert_with(|| Box::new(TrieNode::new()));
                res += node.count;
            }
            node.count += 1;
        }
        res as i32
    }
}