919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
class Node:
    def __init__(self):
        # Array to store links to child nodes,
        # each index represents a letter
        self.links = [None] * 26
        # Counter for number of words
        # that end at this node
        self.cntEndWith = 0
        # Counter for number of words
        # that have this node as a prefix
        self.cntPrefix = 0

    # Check if the node contains
    # a specific key (letter)
    def containsKey(self, ch):
        # Check if the link corresponding
        # to the character exists
        return self.links[ord(ch) - ord('a')] is not None

    # Get the node with a specific
    # key (letter) from the Trie
    def get(self, ch):
        # Return the link
        # corresponding to the character
        return self.links[ord(ch) - ord('a')]

    # Insert a new node with a specific
    # key (letter) into the Trie
    def put(self, ch, node):
        # Set the link corresponding to
        # the character to the provided node
        self.links[ord(ch) - ord('a')] = node

    # Increment the count of words
    # that end at this node
    def increaseEnd(self):
        # Increment the counter
        self.cntEndWith += 1

    # Increment the count of words
    # that have this node as a prefix
    def increasePrefix(self):
        # Increment the counter
        self.cntPrefix += 1

    # Decrement the count of words
    # that end at this node
    def deleteEnd(self):
        # Decrement the counter
        self.cntEndWith -= 1

    # Decrement the count of words
    # that have this node as a prefix
    def reducePrefix(self):
        # Decrement the counter
        self.cntPrefix -= 1

class Trie:
    def __init__(self):
        # Constructor to initialize the
        # Trie with an empty root node
        self.root = Node()

    # Inserts a word into the Trie
    # Time Complexity O(len), where len
    # is the length of the word
    def insert(self, word):
        # Start from the root node
        node = self.root
        # Iterate over each
        # character in the word
        for ch in word:
            if not node.containsKey(ch):
                # Create a new node
                # for the character
                node.put(ch, Node())
            # Move to the child node
            # corresponding to the character
            node = node.get(ch)
            # Increment the prefix
            # count for the node
            node.increasePrefix()
        # Increment the end count
        # for the last node of the word
        node.increaseEnd()

    # Returns the number of words
    # equal to a given word
    def countWordsEqualTo(self, word):
        # Start from the root node
        node = self.root
        # Iterate over each character in the word
        for ch in word:
            if node.containsKey(ch):
                # Move to the child node
                # corresponding to the