2416. 字符串的前缀分数和

题目

给你一个长度为 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] 由小写英文字母组成

解法

我们将 words 存储到字典树中。

在计算分数时，我们计算的是该前缀在多少个 word 中出现，即该路径被复用的次数。

因而我们在插入 word 时，将沿途所有的节点的值都 +1. 节点值即表示从根到当前位置的路径形成的前缀在 words 中出现的次数。

查找时，由于我们计算的是 word 的所有合法前缀的分数之和，因而我们在遍历查找的过程中进行累加操作。

代码


class Trie:
    def __init__(self):
        self.root = {'val': 0}

    def add(self, word):
        node = self.root
        for c in word:
            if c not in node:
                node[c] = {'val': 0}
            node = node[c]
            node['val'] += 1
    
    def search(self, word, i, node):
        c = word[i]
        if c not in node:
            return 0
        node = node[c]
        val = node['val']
        if i < len(word) - 1:
            val += self.search(word, i+1, node)
        return val

class Solution:
    def sumPrefixScores(self, words: List[str]) -> List[int]:
        trie = Trie()
        for word in words:
            trie.add(word)
        ans = [ trie.search(word, 0, trie.root) for word in words ]
        return ans