Implement Trie (Prefix Tree) LeetCode solution in Python

1
class Node:
2
    def __init__(self):
3
        self.child = {}  # to hold the nodes.
4
        self.end = False  # to mark a node if it is the end node or not.
5

6

7
class Trie:
8

9
    def __init__(self):
10
        self.root = Node()
11

12
    def insert(self, word: str) -> None:
13
        # time compl len(word)
14

15
        sz = len(word)
16
        temp = self.root  # to hold the root node.
17

18
        for ind, i in enumerate(word):
19
            if i in temp.child.keys():  # if this curr char in the current node.
20
                temp = temp.child[i]  # another node.
21

22
            else:
23
                temp.child[i] = Node()
24
                temp = temp.child[i]
25

26
            if ind == sz - 1:
27
                temp.end = True
28

29
    def search(self, word: str) -> bool:
30

31
        temp = self.root
32
        for i in word:
33
            if i in temp.child.keys():
34
                temp = temp.child[i]
35
            else:
36
                return 0
37

38
        return temp.end == True
39

40
    def startsWith(self, prefix: str) -> bool:
41
        temp = self.root
42
        for i in prefix:
43
            if i in temp.child.keys():
44
                temp = temp.child[i]
45
            else:
46
                return 0
47
        return 1