prep-coding/leetcode/implement_trie.py at master · tushar-sharma/prep-coding · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
from dataclasses import dataclass, field

@dataclass
class Node:

    __children : dict = field(default_factory=lambda: defaultdict(list))
    __end_word : bool = False

    @property
    def children(self):
        return self.__children

    @property
    def end_word(self) -> bool:
        return self.__end_word

    @end_word.setter
    def end_word(self, end_word) -> None:
        self.__end_word = end_word


class Trie:

    def __init__(self):
        self.__root = Node()


    def insert(self, word: str) -> None:

        current = self.__root

        for c in word:
            children = current.children

            if c not in children:
                children[c] = Node()

            current = children[c]

        current.end_word = True


    def search(self, word: str) -> bool:
        current = self.__root

        for c in word:
            children = current.children

            if c not in children:
                return False

            current = children[c]

        return current.end_word

    def startsWith(self, prefix: str) -> bool:
        current = self.__root

        for c in prefix:
            children = current.children

            if c not in children:
                return False

            current = children[c]

        return True