HackerRank Tries: Contacts

tries_contact.py

class t_node():
    def __init__(self):
        self.children = {}
        self.is_complete = False
        self.valid_children = 0
    
    def add_child(self, char):
        self.children[char] = t_node()
        return self.children[char]
    
    def num_valid_children(self):
        return self.valid_children
        
def add(root, name):
    node = root
    for char in name:
        try:
            node = node.children[char]
            node.valid_children += 1
        except KeyError:
            node = node.add_child(char)
            node.valid_children = 1
    node.is_complete = True
    

def find(root, partial_name):
    node = root
    for char in partial_name:
        try:
            node = node.children[char]
        except KeyError:
            return 0
    return node.num_valid_children()

# should print (char, node) tuples
def test_add():
    name = "joshua"
    root = t_node()
    add(root, name)
    node = root
    for char in name:
        node = node.children[char]
        print(char, node)

# print 3
def test_find():
    root = t_node()
    add(root, "joshua")
    add(root, "jose")
    add(root, "jo")
    print(find(root, "jo"))
        
# Should print 4
def test_t_node():
    root = t_node()
    root.add_child("a")
    a = root.children["a"]
    a.add_child("b")
    a.add_child("c")
    c = a.children["c"]
    c.add_child("d")
    print(root.num_valid_children())


#test_find()
#test_add()
#test_t_node()

if __name__ == "__main__":
    root = t_node()
    n = int(input().strip())
    for a0 in range(n):
        op, contact = input().strip().split(' ')
        if op == "find":
            print(find(root, contact))
        elif op == "add":
            add(root, contact)

"""
# I originally tried to determine
# the number of valid children on each
# find() as below; it was too slow and failed
# ~60% of the test cases with a timeout

def num_valid_children(self):
        valid_children = 0
        if self.is_complete:
            valid_children += 1
        if self.children == {}:
            return valid_children
        else:
            for each in self.children:
                valid_children += self.children[each].num_valid_children()
            return valid_children
"""