leetcode : Sum of Prefix Scores of Strings

https://leetcode.com/problems/sum-of-prefix-scores-of-strings/

 

어제였나? 트라이 공부했던거같은데 이렇게 나와버리네 ㅎㅎ

 

문제를 잘 이해한다면, 그리고 트라이를 알고있다면 쉽게  풀수있다. 하하, 하드문제가 이렇게 쉽게 풀리다니.

기존의 트라이와 같은방법으로 insert를 하되 노드에 prefix_score 변수를 만들어 insert 하면서 한글자씩 타고 들어갈때마다 스코어를 1 증가시켜주면 된다. 즉 insert하면서 스코어를 미리 1씩 증가시켜놓으면 특정 단어를 prefix로 하는 단어가 몇개인지 빠르게 알수있으니까 나중엔 word를 받아 트라이를 쭉 타고 내려가면서 score를 다 더해주기만 하면 sum of prefix scores of strings 완성이다.

 

class Node:
    def __init__(self):
        self.children = {}
        self.prefix_score = 1

class Trie:
    def __init__(self):
        self.root = Node()
        self.root.prefix_score = 0
        
    def insert(self, word):
        node = self.root
        for ch in word:
            if ch not in node.children:
                node.children[ch] = Node()
            else:
                node.children[ch].prefix_score += 1
            node = node.children[ch]
    
    def calc_score(self, word):
        score = 0
        node = self.root
        for ch in word:
            node = node.children[ch]
            score += node.prefix_score
        return score
        

class Solution:
    def sumPrefixScores(self, words: List[str]) -> List[int]:
        trie = Trie()
        for word in words:
            trie.insert(word)
            
        answer = []
        for word in words:
            answer.append(trie.calc_score(word))
            
        return answer