3042. Count Prefix and Suffix Pairs I - Explanation

Description

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.

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.

Constraints:

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

Topics

Company Tags

Please upgrade to NeetCode Pro to view company tags.

Prerequisites

Before attempting this problem, you should be comfortable with:

String Manipulation - Checking prefixes and suffixes of strings using character comparison or built-in methods
Nested Loops - Iterating through all pairs of elements in an array with O(n^2) complexity
Trie Data Structure - Understanding how tries store strings and enable efficient prefix lookups (for optimal solution)

1. Brute Force

Intuition

For each pair of indices (i, j) where i < j, we need to check if words[i] is both a prefix and suffix of words[j]. We compare characters at the beginning and end of words[j] with words[i].

Algorithm

Create a helper function that checks if s1 is both a prefix and suffix of s2.
First verify s1 is not longer than s2.
Compare s1 character by character with the start of s2 (prefix check).
Compare s1 character by character with the end of s2 (suffix check).
Iterate through all pairs (i, j) with i < j and count valid prefix-suffix pairs.

class Solution:
    def countPrefixSuffixPairs(self, words: List[str]) -> int:
        def isPrefixAndSuffix(s1, s2):
            if len(s1) > len(s2):
                return False

            for i in range(len(s1)):
                if s1[i] != s2[i]:
                    return False

            j = 0
            for i in range(len(s2) - len(s1), len(s2)):
                if s1[j] != s2[i]:
                    return False
                j += 1

            return True

        res = 0
        for i in range(len(words)):
            for j in range(i + 1, len(words)):
                res += isPrefixAndSuffix(words[i], words[j])
        return res

public class Solution {
    public int countPrefixSuffixPairs(String[] words) {
        int res = 0;
        for (int i = 0; i < words.length; i++) {
            for (int j = i + 1; j < words.length; j++) {
                if (isPrefixAndSuffix(words[i], words[j])) {
                    res++;
                }
            }
        }
        return res;
    }

    private boolean isPrefixAndSuffix(String s1, String s2) {
        if (s1.length() > s2.length()) return false;

        for (int i = 0; i < s1.length(); i++) {
            if (s1.charAt(i) != s2.charAt(i)) return false;
        }

        int j = 0;
        for (int i = s2.length() - s1.length(); i < s2.length(); i++) {
            if (s1.charAt(j) != s2.charAt(i)) return false;
            j++;
        }

        return true;
    }
}

class Solution {
public:
    int countPrefixSuffixPairs(vector<string>& words) {
        int res = 0;
        for (int i = 0; i < words.size(); i++) {
            for (int j = i + 1; j < words.size(); j++) {
                if (isPrefixAndSuffix(words[i], words[j])) {
                    res++;
                }
            }
        }
        return res;
    }

    bool isPrefixAndSuffix(string s1, string s2) {
        if (s1.size() > s2.size()) return false;

        for (int i = 0; i < s1.size(); i++) {
            if (s1[i] != s2[i]) return false;
        }

        int j = 0;
        for (int i = s2.size() - s1.size(); i < s2.size(); i++) {
            if (s1[j] != s2[i]) return false;
            j++;
        }

        return true;
    }
};

class Solution {
    /**
     * @param {string[]} words
     * @return {number}
     */
    countPrefixSuffixPairs(words) {
        const isPrefixAndSuffix = (s1, s2) => {
            if (s1.length > s2.length) return false;

            for (let i = 0; i < s1.length; i++) {
                if (s1[i] !== s2[i]) return false;
            }

            let j = 0;
            for (let i = s2.length - s1.length; i < s2.length; i++) {
                if (s1[j] !== s2[i]) return false;
                j++;
            }

            return true;
        };

        let res = 0;
        for (let i = 0; i < words.length; i++) {
            for (let j = i + 1; j < words.length; j++) {
                if (isPrefixAndSuffix(words[i], words[j])) {
                    res++;
                }
            }
        }
        return res;
    }
}

public class Solution {
    public int CountPrefixSuffixPairs(string[] words) {
        int res = 0;
        for (int i = 0; i < words.Length; i++) {
            for (int j = i + 1; j < words.Length; j++) {
                if (IsPrefixAndSuffix(words[i], words[j])) {
                    res++;
                }
            }
        }
        return res;
    }

    private bool IsPrefixAndSuffix(string s1, string s2) {
        if (s1.Length > s2.Length) return false;

        for (int i = 0; i < s1.Length; i++) {
            if (s1[i] != s2[i]) return false;
        }

        int j = 0;
        for (int i = s2.Length - s1.Length; i < s2.Length; i++) {
            if (s1[j] != s2[i]) return false;
            j++;
        }

        return true;
    }
}

func countPrefixSuffixPairs(words []string) int {
    isPrefixAndSuffix := func(s1, s2 string) bool {
        if len(s1) > len(s2) {
            return false
        }

        for i := 0; i < len(s1); i++ {
            if s1[i] != s2[i] {
                return false
            }
        }

        j := 0
        for i := len(s2) - len(s1); i < len(s2); i++ {
            if s1[j] != s2[i] {
                return false
            }
            j++
        }

        return true
    }

    res := 0
    for i := 0; i < len(words); i++ {
        for j := i + 1; j < len(words); j++ {
            if isPrefixAndSuffix(words[i], words[j]) {
                res++
            }
        }
    }
    return res
}

class Solution {
    fun countPrefixSuffixPairs(words: Array<String>): Int {
        fun isPrefixAndSuffix(s1: String, s2: String): Boolean {
            if (s1.length > s2.length) return false

            for (i in s1.indices) {
                if (s1[i] != s2[i]) return false
            }

            var j = 0
            for (i in s2.length - s1.length until s2.length) {
                if (s1[j] != s2[i]) return false
                j++
            }

            return true
        }

        var res = 0
        for (i in words.indices) {
            for (j in i + 1 until words.size) {
                if (isPrefixAndSuffix(words[i], words[j])) {
                    res++
                }
            }
        }
        return res
    }
}

class Solution {
    func countPrefixSuffixPairs(_ words: [String]) -> Int {
        func isPrefixAndSuffix(_ s1: String, _ s2: String) -> Bool {
            if s1.count > s2.count { return false }

            let arr1 = Array(s1)
            let arr2 = Array(s2)

            for i in 0..<arr1.count {
                if arr1[i] != arr2[i] { return false }
            }

            var j = 0
            for i in (arr2.count - arr1.count)..<arr2.count {
                if arr1[j] != arr2[i] { return false }
                j += 1
            }

            return true
        }

        var res = 0
        for i in 0..<words.count {
            for j in (i + 1)..<words.count {
                if isPrefixAndSuffix(words[i], words[j]) {
                    res += 1
                }
            }
        }
        return res
    }
}

impl Solution {
    pub fn count_prefix_suffix_pairs(words: Vec<String>) -> i32 {
        let mut res = 0;
        for i in 0..words.len() {
            for j in (i + 1)..words.len() {
                let w1 = &words[i];
                let w2 = &words[j];
                if w2.starts_with(w1.as_str()) && w2.ends_with(w1.as_str()) {
                    res += 1;
                }
            }
        }
        res
    }
}

Time & Space Complexity

Time complexity: $O(n ^ 2 * m)$
Space complexity: $O(1)$

Where $n$ is the size of the input array $words$ , and $m$ is the maximum length of a string.

3. Trie

Intuition

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.

Algorithm

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.

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

class TrieNode {
    Map<String, TrieNode> children = new HashMap<>();
    int count = 0;
}

class Trie {
    TrieNode root = new TrieNode();

    void add(String w) {
        TrieNode cur = root;
        int n = w.length();
        for (int i = 0; i < n; i++) {
            String key = "" + w.charAt(i) + w.charAt(n - 1 - i);
            cur.children.putIfAbsent(key, new TrieNode());
            cur = cur.children.get(key);
            cur.count++;
        }
    }

    int count(String w) {
        TrieNode cur = root;
        int n = w.length();
        for (int i = 0; i < n; i++) {
            String key = "" + w.charAt(i) + w.charAt(n - 1 - i);
            if (!cur.children.containsKey(key)) return 0;
            cur = cur.children.get(key);
        }
        return cur.count;
    }
}

public class Solution {
    public int countPrefixSuffixPairs(String[] words) {
        int res = 0;
        Trie root = new Trie();
        for (int i = words.length - 1; i >= 0; i--) {
            res += root.count(words[i]);
            root.add(words[i]);
        }
        return res;
    }
}

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;
    }
};

class TrieNode {
    constructor() {
        this.children = new Map();
        this.count = 0;
    }
}

class Trie {
    constructor() {
        this.root = new TrieNode();
    }

    /**
     * @param {string} w
     * @return {void}
     */
    add(w) {
        let cur = this.root;
        let n = w.length;
        for (let i = 0; i < n; i++) {
            let key = w[i] + w[n - 1 - i];
            if (!cur.children.has(key)) {
                cur.children.set(key, new TrieNode());
            }
            cur = cur.children.get(key);
            cur.count++;
        }
    }

    /**
     * @param {string} words
     * @return {number}
     */
    count(w) {
        let cur = this.root;
        let n = w.length;
        for (let i = 0; i < n; i++) {
            let key = w[i] + w[n - 1 - i];
            if (!cur.children.has(key)) return 0;
            cur = cur.children.get(key);
        }
        return cur.count;
    }
}

class Solution {
    /**
     * @param {string[]} words
     * @return {number}
     */
    countPrefixSuffixPairs(words) {
        let res = 0;
        let root = new Trie();
        for (let i = words.length - 1; i >= 0; i--) {
            res += root.count(words[i]);
            root.add(words[i]);
        }
        return res;
    }
}

public class TrieNode {
    public Dictionary<string, TrieNode> children = new Dictionary<string, TrieNode>();
    public int count = 0;
}

public class Trie {
    public TrieNode root = new TrieNode();

    public void Add(string w) {
        TrieNode cur = root;
        int n = w.Length;
        for (int i = 0; i < n; i++) {
            string key = w[i].ToString() + w[n - 1 - i];
            if (!cur.children.ContainsKey(key)) {
                cur.children[key] = new TrieNode();
            }
            cur = cur.children[key];
            cur.count++;
        }
    }

    public int Count(string w) {
        TrieNode cur = root;
        int n = w.Length;
        for (int i = 0; i < n; i++) {
            string key = w[i].ToString() + w[n - 1 - i];
            if (!cur.children.ContainsKey(key)) return 0;
            cur = cur.children[key];
        }
        return cur.count;
    }
}

public class Solution {
    public int CountPrefixSuffixPairs(string[] words) {
        int res = 0;
        Trie root = new Trie();
        for (int i = words.Length - 1; i >= 0; i--) {
            res += root.Count(words[i]);
            root.Add(words[i]);
        }
        return res;
    }
}

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

func newTrieNode() *TrieNode {
    return &TrieNode{children: make(map[string]*TrieNode)}
}

type Trie struct {
    root *TrieNode
}

func newTrie() *Trie {
    return &Trie{root: newTrieNode()}
}

func (t *Trie) add(w string) {
    cur := t.root
    n := len(w)
    for i := 0; i < n; i++ {
        key := string(w[i]) + string(w[n-1-i])
        if _, ok := cur.children[key]; !ok {
            cur.children[key] = newTrieNode()
        }
        cur = cur.children[key]
        cur.count++
    }
}

func (t *Trie) getCount(w string) int {
    cur := t.root
    n := len(w)
    for i := 0; i < n; i++ {
        key := string(w[i]) + string(w[n-1-i])
        if _, ok := cur.children[key]; !ok {
            return 0
        }
        cur = cur.children[key]
    }
    return cur.count
}

func countPrefixSuffixPairs(words []string) int {
    res := 0
    trie := newTrie()
    for i := len(words) - 1; i >= 0; i-- {
        res += trie.getCount(words[i])
        trie.add(words[i])
    }
    return res
}

class TrieNode {
    val children = HashMap<String, TrieNode>()
    var count = 0
}

class Trie {
    val root = TrieNode()

    fun add(w: String) {
        var cur = root
        val n = w.length
        for (i in 0 until n) {
            val key = "${w[i]}${w[n - 1 - i]}"
            if (!cur.children.containsKey(key)) {
                cur.children[key] = TrieNode()
            }
            cur = cur.children[key]!!
            cur.count++
        }
    }

    fun count(w: String): Int {
        var cur = root
        val n = w.length
        for (i in 0 until n) {
            val key = "${w[i]}${w[n - 1 - i]}"
            if (!cur.children.containsKey(key)) return 0
            cur = cur.children[key]!!
        }
        return cur.count
    }
}

class Solution {
    fun countPrefixSuffixPairs(words: Array<String>): Int {
        var res = 0
        val trie = Trie()
        for (i in words.size - 1 downTo 0) {
            res += trie.count(words[i])
            trie.add(words[i])
        }
        return res
    }
}

class TrieNode {
    var children = [String: TrieNode]()
    var count = 0
}

class Trie {
    var root = TrieNode()

    func add(_ w: String) {
        var cur = root
        let chars = Array(w)
        let n = chars.count
        for i in 0..<n {
            let key = "\(chars[i])\(chars[n - 1 - i])"
            if cur.children[key] == nil {
                cur.children[key] = TrieNode()
            }
            cur = cur.children[key]!
            cur.count += 1
        }
    }

    func getCount(_ w: String) -> Int {
        var cur = root
        let chars = Array(w)
        let n = chars.count
        for i in 0..<n {
            let key = "\(chars[i])\(chars[n - 1 - i])"
            if cur.children[key] == nil { return 0 }
            cur = cur.children[key]!
        }
        return cur.count
    }
}

class Solution {
    func countPrefixSuffixPairs(_ words: [String]) -> Int {
        var res = 0
        let trie = Trie()
        for i in stride(from: words.count - 1, through: 0, by: -1) {
            res += trie.getCount(words[i])
            trie.add(words[i])
        }
        return res
    }
}

struct TrieNode {
    children: HashMap<(u8, u8), TrieNode>,
    count: i32,
}

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

struct Trie {
    root: TrieNode,
}

impl Trie {
    fn new() -> Self {
        Trie { root: TrieNode::new() }
    }

    fn add(&mut self, w: &[u8]) {
        let mut cur = &mut self.root;
        let n = w.len();
        for i in 0..n {
            let key = (w[i], w[n - 1 - i]);
            cur = cur.children.entry(key).or_insert_with(TrieNode::new);
            cur.count += 1;
        }
    }

    fn count(&self, w: &[u8]) -> i32 {
        let mut cur = &self.root;
        let n = w.len();
        for i in 0..n {
            let key = (w[i], w[n - 1 - i]);
            match cur.children.get(&key) {
                Some(node) => cur = node,
                None => return 0,
            }
        }
        cur.count
    }
}

impl Solution {
    pub fn count_prefix_suffix_pairs(words: Vec<String>) -> i32 {
        let mut res = 0;
        let mut trie = Trie::new();
        for i in (0..words.len()).rev() {
            let w = words[i].as_bytes();
            res += trie.count(w);
            trie.add(w);
        }
        res
    }
}

Time & Space Complexity

Time complexity: $O(n * m)$
Space complexity: $O(n * m)$

Where $n$ is the size of the input array $words$ , and $m$ is the maximum length of a string.

Common Pitfalls

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.

# Wrong: only checks prefix
if w2.startswith(w1):
    res += 1
# Correct: checks both conditions
if w2.startswith(w1) and w2.endswith(w1):
    res += 1

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.