【前缀树】字符串的前缀分数和【力扣第 311 场周赛】

题目

6183. 字符串的前缀分数和

---

给你一个长度为 n 的数组 words ，该数组由 非空 字符串组成。

定义字符串 word 的 分数 等于以 word 作为 前缀 的 words[i] 的数目。

• 例如，如果 words = ["a", "ab", "abc", "cab"] ，那么 "ab" 的分数是 2 ，因为 "ab" 是 "ab" 和 "abc" 的一个前缀。

返回一个长度为 n 的数组 answer ，其中 answer[i] 是 words[i] 的每个非空前缀的分数 总和 。

**注意：**字符串视作它自身的一个前缀。

示例 1：

输入：words = ["abc","ab","bc","b"]
输出：[5,4,3,2]
解释：对应每个字符串的答案如下：
- "abc" 有 3 个前缀："a"、"ab" 和 "abc" 。
- 2 个字符串的前缀为 "a" ，2 个字符串的前缀为 "ab" ，1 个字符串的前缀为 "abc" 。
总计 answer[0] = 2 + 2 + 1 = 5 。
- "ab" 有 2 个前缀："a" 和 "ab" 。
- 2 个字符串的前缀为 "a" ，2 个字符串的前缀为 "ab" 。
总计 answer[1] = 2 + 2 = 4 。
- "bc" 有 2 个前缀："b" 和 "bc" 。
- 2 个字符串的前缀为 "b" ，1 个字符串的前缀为 "bc" 。 
总计 answer[2] = 2 + 1 = 3 。
- "b" 有 1 个前缀："b"。
- 2 个字符串的前缀为 "b" 。
总计 answer[3] = 2 。

示例 2：

输入：words = ["abcd"]
输出：[4]
解释：
"abcd" 有 4 个前缀 "a"、"ab"、"abc" 和 "abcd"。
每个前缀的分数都是 1 ，总计 answer[0] = 1 + 1 + 1 + 1 = 4 。

提示：

• 1 <= words.length <= 1000
• 1 <= words[i].length <= 1000
• words[i] 由小写英文字母组成

解题

方法一：前缀树(Trie)

思路

前缀树的原理及实现：【前缀树】实现 Trie (前缀树)

把数组中的所有字符串存入字典树，由于每个节点都是其子树节点的前缀，题目中所说的字符串分数就是经过该节点的字符串个数，即以该节点为前缀的字符串的个数。

所以可以在前缀树插入的同时再记录每个节点的所表示字符出现的次数(count)，然后在查找前缀时，把找出的路径中节点的 count 累加起来就是该字符串的前缀分数和。

代码

class Solution {
    public int[] sumPrefixScores(String[] words) {
        Trie trie = new Trie();
        for (String word : words) trie.insert(word);
        int n = words.length;
        int[] ans = new int[n];
        for (int i = 0; i < n; ++i) {
            ans[i] = trie.searchPrefix(words[i]);
        }
        return ans;
    }
    
    static class Trie {
        private boolean isEnd;
        private Trie[] children;
        private int count;

        public Trie() {
            isEnd = false;
            children = new Trie[26];
            count = 0;
        }

        public void insert(String word) {
            Trie node = this;
            for (char ch : word.toCharArray()) {
                int idx = ch - 'a';
                if (node.children[idx] == null) node.children[idx] = new Trie();
                ++node.children[idx].count;
                node = node.children[idx];
            }
            node.isEnd = true;
        }

        private int searchPrefix(String prefix) {
            Trie node = this;
            int cnt = 0;
            for (char ch : prefix.toCharArray()) {
                int idx = ch - 'a';
                if (node.children[idx] == null) return 0;
                node = node.children[idx];
                cnt += node.count;
            }
            return node == null ? 0 : cnt;
        }
    }
}