Solution for HackerRank challenge "Contacts" (https://www.hackerrank.com/challenges/contacts/problem)

hackerrank-contacts.py

#!/bin/python3

DEBUG = False

import os
import sys

def trie_find(node, path):
    if not node.get(path[0]):
        return 0

    if len(path) > 1:
        return trie_find(node[path[0]], path[1:])
    else:
        return node[path[0]]['__count']

def trie_inc(node, path):
    if not node.get(path[0]):
        node[path[0]] = { '__count': 0 }

    node[path[0]]['__count'] = node[path[0]]['__count'] + 1

    if len(path) > 1:
        trie_inc(node[path[0]], path[1:])
        

#
# Complete the contacts function below.
#
def contacts(queries):
    trie = {}
    answers = []

    for query in queries:
        if query[0] == 'add':
            trie_inc(trie, list(query[1]))
        else:
            answers.append(trie_find(trie, list(query[1])))

    return answers

if __name__ == '__main__':
    if DEBUG:
        queries = [line.strip().split() for line in open("in.txt", 'r')][1:]
        
        result = contacts(queries)

        for line in result:
            print(line)
    else:
        fptr = open(os.environ['OUTPUT_PATH'], 'w')

        queries_rows = int(input())

        queries = []

        for _ in range(queries_rows):
            queries.append(input().rstrip().split())

        result = contacts(queries)

        fptr.write('\n'.join(map(str, result)))
        fptr.write('\n')

        fptr.close()