hackerrank solution - cut the tree

cutTheTree.py

#!/bin/python3

import math
import os
import random
import re
import sys


#
# Complete the 'cutTheTree' function below.
#
# The function is expected to return an INTEGER.
# The function accepts following parameters:
#  1. INTEGER_ARRAY data
#  2. 2D_INTEGER_ARRAY edges
#


def dfs(node, par, adjlst, K):
    conlst = [0] * (K + 1)
    dislst = [0] * (K + 1)
    conlst[0] = 1
    for chld in adjlst[node]:
        if chld != par:
            tcl, tdl = dfs(chld, node, adjlst, K)
            # print(str(tcl))
            # print(str(tdl))
            dislst[0] += tdl[0]
            tcl2 = [0] * (K + 1)
            for i in range(K):
                dislst[i + 1] += tcl[i] + tdl[i + 1]
                tcl2[i + 1] += conlst[i]
            for i in range(K + 1):
                for j in range(K + 1):
                    if i + j <= K:
                        tcl2[i + j] += tcl[i] * conlst[j]
            conlst = list(tcl2)
    return conlst, dislst



def cutTheTree(data, edges):
    #
    n = int(line[0])
    K = int(line[1])

    adjlst = [[1]]
    for i in range(n):
        adjlst.append([])
    adjlst[1].append(0)
    
    for i in range(n - 1):
        x = int(line[0])
        y = int(line[1])
        adjlst[x].append(y)
        adjlst[y].append(x)    
    # Write your code here
    conlst, dislst = dfs(1, 0, adjlst, K)
    # print(str(conlst))
    # print(str(dislst))
    return sum(conlst) + sum(dislst) + 1
    

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    n = int(input().strip())

    data = list(map(int, input().rstrip().split()))

    edges = []

    for _ in range(n - 1):
        edges.append(list(map(int, input().rstrip().split())))

    result = cutTheTree(data, edges)

    fptr.write(str(result) + '\n')

    fptr.close()