Problem statement

You are given an array of lowercase strings words. Say that isPrefixAndSuffix(str1, str2) is true when str1 is both a prefix and a suffix of str2 — for example "co" and "coco", but not "co" and "code".

Return the number of index pairs (i, j) with i < j for which isPrefixAndSuffix(words[i], words[j]) is true.

CodeKairo bound: the original allows 10^5 words, whose answer needs 64 bits; here words.length <= 5 * 10^4, so the answer fits a 32-bit integer.

Example 1

Input: words = ["co","code","coco","c"]
Output: 1
Explanation: Only `(0, 2)`: "co" starts and ends "coco". "code" starts with "co" but does not end with it.

Example 2

Input: words = ["a","a","a"]
Output: 3

Example 3

Input: words = ["abc","c","cbc","cabc"]
Output: 2
Explanation: `(1, 2)` and `(1, 3)`.

Constraints

• 1 <= words.length <= 5 * 10^4
• 1 <= words[i].length <= 10^5
• The total length of all words is at most 5 * 10^5
• words[i] consists only of lowercase English letters

Hints for Count Prefix and Suffix Pairs II

1. Hint 1

   a is a prefix and a suffix of b exactly when, for every k < a.length, a[k] == b[k] and a[len(a)-1-k] == b[len(b)-1-k].

2. Hint 2

   So read each word as a sequence of character pairs (w[k], w[len-1-k]). The condition becomes "the pair sequence of a is a prefix of the pair sequence of b".

3. Hint 3

   Insert the words, in order, into a trie over these pairs. While walking word j, every earlier word that ended at a node on the path is a match — add those end counts.

How to solve Count Prefix and Suffix Pairs II

Pair each position with its mirror: word w becomes the sequence (w[0], w[L-1]), (w[1], w[L-2]), …. Being a prefix and a suffix at once is then just being a prefix of that pair sequence, which a trie counts directly.

Approach

1. Keep a trie whose edges are labelled by a pair of letters (676 possible labels, stored in a hash map), and an ends counter per node.
2. For each word s in index order, walk the trie along its pairs (s[k], s[L-1-k]) for k = 0..L-1, creating missing nodes.
3. At every node reached, add ends[node] to the answer — those are earlier words whose whole pair sequence matches this prefix.
4. After the walk, increment ends of the final node.

Why it works

An earlier word a of length m ends at the node of depth m on s's path exactly when its first m pairs equal s's first m pairs, i.e. a[k] = s[k] (prefix) and a[m-1-k] = s[L-1-k] (suffix) for all k < m. Querying before inserting s counts only pairs with i < j, each once.

Complexity

• Time — O(total length)
• Space — O(total length)

Pitfalls

• Checking prefix and suffix separately with two tries does not work — the two matches must come from the same earlier word.
• Equal words count: "a" is a prefix and a suffix of "a".
• Add ends at every node on the path, not only at the last one.

Reference solution

Python

from typing import List

def countPrefixSuffixPairs(words: List[str]) -> int:
    child = {}
    ends = [0]
    total = 0
    for s in words:
        length = len(s)
        node = 0
        for k in range(length):
            key = node * 676 + (ord(s[k]) - 97) * 26 + (ord(s[length - 1 - k]) - 97)
            nxt = child.get(key)
            if nxt is None:
                nxt = len(ends)
                ends.append(0)
                child[key] = nxt
            node = nxt
            total += ends[node]
        ends[node] += 1
    return total

Java

public static int countPrefixSuffixPairs(String[] words) {
    int cap = 1;
    for (String s : words) cap += s.length();
    int[] ends = new int[cap];
    HashMap<Integer, Integer> child = new HashMap<>();
    int nodes = 1;
    long total = 0;
    for (String s : words) {
        int len = s.length(), node = 0;
        for (int k = 0; k < len; k++) {
            int key = node * 676 + (s.charAt(k) - 'a') * 26 + (s.charAt(len - 1 - k) - 'a');
            Integer nxt = child.get(key);
            if (nxt == null) {
                nxt = nodes++;
                child.put(key, nxt);
            }
            node = nxt;
            total += ends[node];
        }
        ends[node]++;
    }
    return (int) total;
}

C++

int countPrefixSuffixPairs(vector<string>& words) {
    unordered_map<int, int> child;
    vector<int> ends(1, 0);
    long long total = 0;
    for (const string& s : words) {
        int len = s.size(), node = 0;
        for (int k = 0; k < len; k++) {
            int key = node * 676 + (s[k] - 'a') * 26 + (s[len - 1 - k] - 'a');
            auto it = child.find(key);
            int nxt;
            if (it == child.end()) {
                nxt = ends.size();
                ends.push_back(0);
                child[key] = nxt;
            } else {
                nxt = it->second;
            }
            node = nxt;
            total += ends[node];
        }
        ends[node]++;
    }
    return (int) total;
}

JavaScript

var countPrefixSuffixPairs = function(words) {
    var child = new Map();
    var ends = [0];
    var total = 0;
    for (var w = 0; w < words.length; w++) {
        var s = words[w], len = s.length, node = 0;
        for (var k = 0; k < len; k++) {
            var key = node * 676 + (s.charCodeAt(k) - 97) * 26 + (s.charCodeAt(len - 1 - k) - 97);
            var nxt = child.get(key);
            if (nxt === undefined) {
                nxt = ends.length;
                ends.push(0);
                child.set(key, nxt);
            }
            node = nxt;
            total += ends[node];
        }
        ends[node]++;
    }
    return total;
};

Also on the editorial tab: C, C#, Go, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

More arrays problems

• Design Movie Rental System Hard
• Prefix and Suffix Search Hard
• Height of Binary Tree After Subtree Removal Queries Hard
• Stamping the Grid Hard
• Maximum Number of Ways to Partition an Array Hard
• Smallest Rotation with Highest Score Hard