Trie数据结构前缀词计数功能实现及代码优化问询

Hey there! Let's fix your Trie implementation to add the prefix count feature, plus clean up some parts for better efficiency and clarity.

First, the core missing piece in your current code is a way to track how many words pass through each node. We'll add a prefix_count attribute to each node for this purpose. We'll also fix small issues to make the Trie more robust.

Modified & Optimized Code #

class TrieNode:
    def __init__(self):
        # Initialize 26 slots for each lowercase letter
        self.children = [None] * 26
        # Tracks how many words use this node as part of their prefix
        self.prefix_count = 0
        # Marks if this node is the end of a complete word
        self.is_end_of_word = False

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

    def _get_char_index(self, char):
        # Convert char to 0-25 index (handles uppercase by converting to lowercase first)
        return ord(char.lower()) - ord('a')

    def insert(self, word):
        current_node = self.root
        for char in word:
            index = self._get_char_index(char)
            if not current_node.children[index]:
                current_node.children[index] = TrieNode()
            # Move to the child node and increment its prefix count
            current_node = current_node.children[index]
            current_node.prefix_count += 1
        # Mark the end of the inserted word
        current_node.is_end_of_word = True

    def search(self, word):
        # Now properly checks if the complete word exists (not just the prefix)
        current_node = self.root
        for char in word:
            index = self._get_char_index(char)
            if not current_node.children[index]:
                return False
            current_node = current_node.children[index]
        return current_node.is_end_of_word

    def count_prefix(self, prefix):
        # Returns the number of words starting with the given prefix
        current_node = self.root
        for char in prefix:
            index = self._get_char_index(char)
            if not current_node.children[index]:
                # No words match this prefix at all
                return 0
            current_node = current_node.children[index]
        # The prefix_count here is the total number of words using this prefix
        return current_node.prefix_count

# Test your example scenario
if __name__ == "__main__":
    trie = Trie()
    words = ['Stack', 'Stackoverflow', 'Stackover', 'LOL']
    for word in words:
        trie.insert(word)
    
    print(trie.count_prefix('Stack'))  # Output: 3 (matches your expected result)
    print(trie.search('Stack'))        # Output: True (confirms 'Stack' is a stored word)
    print(trie.search('Stackov'))      # Output: False (correctly identifies it's not a complete word)

Key Changes & Explanations #

• prefix_count Attribute: Every time we insert a word, we increment the prefix_count of each node we traverse. When we reach the end of the target prefix, this count directly tells us how many words start with that prefix.
• is_end_of_word Flag: Fixes your original search method—now it correctly checks if a complete word exists, not just if the prefix path exists.
• Simplified Character Indexing: Replaced your manual letter-to-index dictionary with ord(char.lower()) - ord('a')—this is more efficient, cleaner, and handles uppercase letters automatically.
• Clear Naming: Renamed node to TrieNode for readability, and removed the redundant val attribute (since we can derive the character from the index if needed).

Additional Optimizations to Consider #

• Case Sensitivity: If you want strict case sensitivity, just remove the .lower() call in _get_char_index.
• Memory Efficiency: For sparse tries (lots of unused child slots), you could replace the list of children with a dictionary (e.g., self.children = {}) to save space, though list access is faster for 26 fixed slots.
• Delete Functionality: If you ever need to remove words, you'd decrement prefix_count during deletion and clean up empty nodes if needed.

内容的提问来源于stack exchange，提问作者Tanmay Shrivastava