Cod sursa(job #3169250)

Utilizator vali2wdAdrian Valentin Nafornita vali2wd Data 14 noiembrie 2023 15:29:49
Problema Trie Scor 25
Compilator py Status done
Runda Arhiva educationala Marime 2.92 kb
# https://www.infoarena.ro/problema/trie


# 0 w - adauga o aparitie a cuvantului w in lista
# 1 w - sterge o aparitie a cuvantului w din lista
# 2 w - tipareste numarul de aparitii ale cuvantului w in lista
# 3 w - tipareste lungimea celui mai lung prefix comun al lui w cu oricare cuvant din lista

# redirect to file all prints
import sys
import gc
sys.stdout = open("trie.out","w")

class TrieNode:
    def __init__(self):
        self.children = [[chr(ord('a') + i), 0, None] for i in range(26)] # lista de copii
        self.count = 0 # pentru cate cuvinte se termina in nodul curent


def ins(node: TrieNode, word: str):
    if len(word) == 0:
        node.count += 1
        return

    index = ord(word[0]) - ord('a')
    node.children[index][1] += 1 # reprezinta numarul de copii care urmeaza sa se formeze din nodul curent
    if node.children[index][2] is None:
        node.children[index][2] = TrieNode()

    ins(node.children[index][2], word[1:])


def pre(node: TrieNode, word:str, length = 0):
        if len(word) == 0:
            return length
        # if node.children[ord(word[0]) - ord('a')][1] == 0 and node.count != 0:
        #     return length + 1

        if node.children[ord(word[0]) - ord('a')][1] > 0 and node.children[ord(word[0]) - ord('a')][2] is not None:
            return pre(node.children[ord(word[0]) - ord('a')][2], word[1:], length + 1)
        return length

def occurence(node: TrieNode, word:str):
    if len(word) == 0:
        return node.count
    if node.children[ord(word[0]) - ord('a')][2] is None:
        return 0
    return occurence(node.children[ord(word[0]) - ord('a')][2], word[1:])

def delete(node: TrieNode, word: str):
    if len(word) == 0:
        if node.count > 0:
            node.count -= 1
        gc.collect()
        return

    index = ord(word[0]) - ord('a')
    node.children[index][1] -= 1

    if node.children[index][2] is None:
        return



    delete(node.children[index][2], word[1:])

    if node.children[index][2] is not None and node.children[index][1] == 0:
        node.children[index][2] = None



# driver code:
if __name__ == "__main__":
    trie = TrieNode()
    file = open("trie.in", "r")

    # ins(trie, "ana")
    # ins(trie, "ana")
    # ins(trie, "ananas")
    # ins(trie, "banane")
    # print(occurence(trie, "ana"))
    # myana = "ana"
    # delete(trie, myana)
    # print(occurence(trie, "ana"))
    # delete(trie, myana)
    # print(occurence(trie, "ana"))
    # delete(trie, myana)
    # print(occurence(trie, "ana"))

    for line in file:
        command, word = line.split()
        if command == "0":
            ins(trie, word)
        elif command == "1":
            delete(trie, word)
        elif command == "2":
            print(occurence(trie,word))
        elif command == "3":
            print(pre(trie, word))
        else:
            print("Invalid command")