diff --git a/ch6/Trie.py b/ch6/Trie.py index bfd0170..dea94eb 100644 --- a/ch6/Trie.py +++ b/ch6/Trie.py @@ -1,53 +1,62 @@ -class vertex: - def __init__(self, chr: str): - self.alphabet = chr - self.exist = False - self.child = [None] * 26 +def norm_codepoint(s: str, base: int = ord("A")) -> int: + """Return integer represening normalized codepoint for a given base.""" + return ord(s)-base + +class Vertex: + def __init__(self, ch: str): + self.alphabet = ch + self.exist = False + self.child = [None] * 26 class Trie: - def __init__(self): - self.root = vertex('!') - - def insert(self, word): - cur = self.root - for w in word: - alphaNum = ord(w)-ord('A') - if cur.child[alphaNum] == None: - cur.child[alphaNum] = vertex(w) - cur = cur.child[alphaNum] - cur.exist = True - - def search(self, word): - cur = self.root - for w in word: - alphaNum = ord(w)-ord('A') - if cur.child[alphaNum] == None: - return False - cur = cur.child[alphaNum] - return cur.exist - - def startsWith(self, prefix): - cur = self.root - for w in prefix: - alphaNum = ord(w)-ord('A') - if cur.child[alphaNum] == None: - return False - cur = cur.child[alphaNum] - return True - - -def main(): - T = Trie() - S = ['CAR', 'CAT', 'RAT'] - for s in S: - print('Insert', s) - T.insert(s) - - print('\'CAR\' exist?', T.search("CAR")) - print('\'DOG\' exist?', T.search("DOG")) - print('Starts with \'CA\' exist?', T.startsWith("CA")) - print('Starts with \'Z\' exist?', T.startsWith("Z")) - print('Starts with \'AT\' exist?', T.startsWith("AT")) - - -main() + def __init__(self): + self.root = Vertex("!") + + def insert(self, word: str) -> None: + cur = self.root + for char in word: + alpha_num = norm_codepoint(char) + current_child = cur.child[alpha_num] + # If current_child is None, set cur as new Vertex instance. + cur = current_child or Vertex(char) + cur.exist = True + + def search(self, word: str) -> bool: + cur = self.root + for char in word: + alpha_num = norm_codepoint(char) + current_child = cur.child[alpha_num] + if not current_child: + return False + cur = current_child + return cur.exist + + def starts_with(self, prefix: str) -> bool: + cur = self.root + for char in prefix: + alpha_num = norm_codepoint(char) + current_child = cur.child[alpha_num] + if not current_child: + return False + cur = current_child + return True + + +if __name__ == "__main__": + t = Trie() + words = "CAR", "CAT", "RAT" + for word in words: + print("Inserting: ", word) + t.insert(word) + + print() + + prefixes = "CA", "Z", "AT" + results = map(t.starts_with, prefixes) + for prefix, result in zip(prefixes, results): + print(f"""Starts with "{prefix}" exist? {result}""") + + print() + + print(f""""CAR" exist? {t.search(words[0])}""") + print(f""""DOG" exist? {t.search("DOG")}""")