题目描述

给你一个下标从 0 开始的字符串数组 words。

定义一个 布尔 函数 isPrefixAndSuffix，它接受两个字符串参数 str1 和 str2：

• 当 str1 同时是 str2 的前缀（prefix）和后缀（suffix）时，isPrefixAndSuffix(str1, str2) 返回 true，否则返回 false。

例如，isPrefixAndSuffix("aba", "ababa") 返回 true，因为 "aba" 既是 "ababa" 的前缀，也是 "ababa" 的后缀，但是 isPrefixAndSuffix("abc", "abcd") 返回 false。

以整数形式，返回满足 i < j 且 isPrefixAndSuffix(words[i], words[j]) 为 true 的下标对 (i, j) 的 数量。

样例

输入：words = ["a","aba","ababa","aa"]
输出：4
解释：在本示例中，计数的下标对包括：
i = 0 且 j = 1，因为 isPrefixAndSuffix("a", "aba") 为 true。
i = 0 且 j = 2，因为 isPrefixAndSuffix("a", "ababa") 为 true。
i = 0 且 j = 3，因为 isPrefixAndSuffix("a", "aa") 为 true。
i = 1 且 j = 2，因为 isPrefixAndSuffix("aba", "ababa") 为 true。
因此，答案是 4。

输入：words = ["pa","papa","ma","mama"]
输出：2
解释：在本示例中，计数的下标对包括：
i = 0 且 j = 1，因为 isPrefixAndSuffix("pa", "papa") 为 true。
i = 2 且 j = 3，因为 isPrefixAndSuffix("ma", "mama") 为 true。
因此，答案是 2。

输入：words = ["abab","ab"]
输出：0
解释：在本示例中，唯一有效的下标对是 i = 0 且 j = 1，
    但是 isPrefixAndSuffix("abab", "ab") 为 false。
因此，答案是 0。

限制

• 1 <= words.length <= 10^5
• 1 <= words[i].length <= 10^5
• words[i] 仅由小写英文字母组成。
• 所有 words[i] 的长度之和不超过 5 * 10^5。

算法

(字典树，KMP) $O(|\Gamma| \cdot \sum |words_i|)$

1. 考虑使用字典树，从后往前遍历每个字符串，先查询匹配这个字符串，然后将这个字符串插入到字典树中。
2. 插入到字典树时，需要先对这个字符串求 KMP，目的是找到所有前缀等于后缀的位置，并将这些位置对应字典树节点上的权值自增 1。
3. 查询时，在字典树上匹配当前字符串，并返回最终匹配节点的权值。
4. 在代码实现中，将匹配和插入的过程写到了一起。

时间复杂度

• 设字符集的大小为 $|\Gamma|$，则构造字典树的总时间复杂度为 $O(|\Gamma| \cdot \sum |words_i|)$。
• 每个字符串求 KMP 的时间复杂度为 $O(\sum |words_i|)$)。
• 匹配查询的时间复杂度为 $O(\sum |words_i|)$)。
• 故总时间复杂度为 $O(|\Gamma| \cdot \sum |words_i|)$。

空间复杂度

• 需要 $O(|\Gamma| \cdot \sum |words_i|)$ 的额外空间存储字典树。

C++ 代码

#define LL long long

const int N = 100010;

struct Node {
    Node *nxt[26];
    int cnt;

    Node() {
        memset(nxt, 0, sizeof(nxt));
        cnt = 0;
    }
};

class Solution {
private:
    Node *r, *idx[N];
    int ne[N];

    int solve(const string &s) {
        const int n = s.size();

        ne[0] = -1;
        for (int i = 1, j = -1; i < n; i++) {
            while (j > -1 && s[j + 1] != s[i])
                j = ne[j];

            if (s[j + 1] == s[i])
                ++j;

            ne[i] = j;
        }

        Node *p = r;
        for (int i = 0; i < n; i++) {
            int t = s[i] - 'a';

            if (p->nxt[t] == NULL)
                p->nxt[t] = new Node();

            p = p->nxt[t];
            idx[i] = p;
        }

        int res = p->cnt;

        for (int i = n - 1; i > -1; i = ne[i])
            ++idx[i]->cnt;

        return res;
    }

public:
    LL countPrefixSuffixPairs(vector<string>& words) {
        r = new Node();

        LL ans = 0;
        for (int i = words.size() - 1; i >= 0; i--)
            ans += solve(words[i]);

        return ans;
    }
};