2416. Sum of Prefix Scores of Strings

Description

The problem involves calculating the prefix score for each word in a given list of words. The prefix score for a word is the sum of counts of prefixes in all words in the list. This problem requires the efficient construction of a prefix-based data structure to quickly calculate prefix counts for each word.

Intuition

The intuition behind the solution is to recognize that the problem revolves around prefixes of words. A Trie (prefix tree) is well-suited for storing words in such a way that common prefixes can be counted efficiently. For each word, we can traverse the Trie and compute the total number of times each prefix has appeared among all the words.

Approach

1. Build a Trie: Start by inserting each word from the list into a Trie. As each character of a word is inserted, we update the count of how many times that prefix has been encountered.
2. Prefix Score Calculation: For each word, traverse through the Trie and accumulate the prefix count. Each prefix encountered contributes to the score.
3. Return the Results: For every word, store its prefix score in a result list and return this list at the end.

Steps:

• Construct a Trie data structure that stores the number of words that share a given prefix.
• Insert all words into the Trie.
• For each word, traverse through the Trie and sum up the prefix counts to calculate the prefix score.
• Return the prefix score for all words.

Complexity

Time Complexity:

• Insertion of each word into the Trie takes (O(L)), where (L) is the length of the word. Therefore, inserting all words takes (O(N \cdot L)), where (N) is the number of words.
• Calculating the prefix score for each word also takes (O(L)). Therefore, the total time complexity is (O(N \cdot L)).

Space Complexity:

• The space complexity is (O(N \cdot L)) due to the space required to store the Trie, where (N) is the number of words and (L) is the average length of each word.

Code

C++

class TrieNode {
public:
    unordered_map children;
    int prefixCount;
    
    TrieNode() {
        prefixCount = 0;
    }
};

class Trie {
public:
    TrieNode* root;
    
    Trie() {
        root = new TrieNode();
    }
    
    void insert(const string& word) {
        TrieNode* node = root;
        for (char c : word) {
            if (node->children.find(c) == node->children.end()) {
                node->children[c] = new TrieNode();
            }
            node = node->children[c];
            node->prefixCount++;
        }
    }
    
    int getPrefixScore(const string& word) {
        TrieNode* node = root;
        int score = 0;
        for (char c : word) {
            node = node->children[c];
            score += node->prefixCount;
        }
        return score;
    }
};

class Solution {
public:
    vector sumPrefixScores(vector& words) {
        Trie trie;
        vector result;
        
        // Insert all words into the trie
        for (const string& word : words) {
            trie.insert(word);
        }
        
        // Get the prefix score for each word
        for (const string& word : words) {
            result.push_back(trie.getPrefixScore(word));
        }
        
        return result;
    }
};

Python

class TrieNode:
    def __init__(self):
        self.children = {}
        self.prefixCount = 0

class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
            node.prefixCount += 1

    def getPrefixScore(self, word):
        node = self.root
        score = 0
        for char in word:
            node = node.children[char]
            score += node.prefixCount
        return score

class Solution:
    def sumPrefixScores(self, words):
        trie = Trie()
        result = []
        for word in words:
            trie.insert(word)
        for word in words:
            result.append(trie.getPrefixScore(word))
        return result

Java

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

class Trie {
    TrieNode root;

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

    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            node.children.putIfAbsent(ch, new TrieNode());
            node = node.children.get(ch);
            node.prefixCount++;
        }
    }

    public int getPrefixScore(String word) {
        TrieNode node = root;
        int score = 0;
        for (char ch : word.toCharArray()) {
            node = node.children.get(ch);
            score += node.prefixCount;
        }
        return score;
    }
}

public class Solution {
    public int[] sumPrefixScores(String[] words) {
        Trie trie = new Trie();
        int[] result = new int[words.length];
        for (String word : words) {
            trie.insert(word);
        }
        for (int i = 0; i < words.length; i++) {
            result[i] = trie.getPrefixScore(words[i]);
        }
        return result;
    }
}

JavaScript

/**
 * @param {string[]} words
 * @return {number[]}
 */
class TrieNode {
    constructor() {
        this.children = {}; //a trie
        this.prefixCount = 0; 
    }
}
class Trie {
    constructor() {
        this.root = new TrieNode();
    }
    insert(word) {
        let node = this.root;
        for (const char of word) {
            if (!node.children[char]) { //Create a Trie if not exists
                node.children[char] = new TrieNode();
            }
            node = node.children[char]; //add element to the Trie
            node.prefixCount++;
        }
    }

    getPrefixScore(word) {
        let node = this.root;
        let score = 0;
        for (const char of word) {
            node = node.children[char];
            score += node.prefixCount;
        }
        return score;
    }
}
var sumPrefixScores = function(words) {
    const trie = new Trie();
    for (const word of words) {
        trie.insert(word);
    }
    const result = words.map(word => trie.getPrefixScore(word));
    
    return result;
};