Problem :

https://leetcode.com/problems/implement-trie-prefix-tree/


My Solution :

class TrieNode:

def __init__(self):
self.children = [None]*26
self.is_end = False


class Trie:

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

def insert(self, word: str) -> None:
curr = self.root
for c in word:
idx = ord(c)-97
curr.children[idx] = curr.children[idx] or TrieNode()
curr = curr.children[idx]
curr.is_end = True

def search(self, word: str, is_startswith=False) -> bool:
curr = self.root
for c in word:
idx = ord(c)-97
curr = curr.children[idx]
if not curr:
return False
return is_startswith or curr.is_end

def startsWith(self, prefix: str) -> bool:
return self.search(prefix, True)


Comment :

위는 알파벳 소문자라는 제약이 있어서 크기 26짜리 list를 사용한 것이고, 아래는 아무 문자에 대응하기 위해 dictionary를 사용한 것.


My Solution 2 :

class TrieNode:

def __init__(self):
self.children = dict()
self.is_end = False


class Trie:

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

def insert(self, word: str) -> None:
curr = self.root
for c in word:
curr = curr.children.setdefault(c, TrieNode())
curr.is_end = True

def search(self, word: str, is_startswith=False) -> bool:
curr = self.root
for c in word:
curr = curr.children.get(c)
if not curr:
return False
return is_startswith or curr.is_end

def startsWith(self, prefix: str) -> bool:
return self.search(prefix, True)