TrieMap implementation in python

triemap.py

class Node:
  def __init__(self, val = None):
    self.left = None
    self.right = None
    self.val = val
    self.end = False

class TrieMap:
  def __init__(self):
    self.root = Node()
    
  def set(self, key, val):
    curr = self.root
    while key:
      if key & 1:
        if curr.right == None:
          curr.right = Node()
        curr = curr.right
      else:
        if curr.left == None:
          curr.left = Node()
        curr = curr.left
      key = key >> 1
    curr.val = val
    curr.end = True
    
  def get(self, key):
    curr = self.root
    while key:
      if curr == None:
        return None
      if key & 1:
        curr = curr.right
      else:
        curr = curr.left
      key = key >> 1
    if not curr or not curr.end:
      return None
    return curr.val
    
tmap = TrieMap()
tmap.set(41, "Hello World")
print(tmap.get(41))