TLE in python

I am getting TLE only in (subtask 4,task 10) on running following code. I have got 60/100 score. Please help.

from sys import setrecursionlimit
setrecursionlimit(10**8)

def dfs(root, depth, parent = -1):
sp = -1
parent_list[root] = parent
for e in adj_list[root]:
if e != parent_list[root]:
temp = dfs(e, depth + 1, root)
if temp != -1 and sp == -1:
sp = temp

if sp_nodes[root]:
    sp = root
if sp != -1:
    ans_list[root] = depth
    ans_node_list[root] = sp

return sp

def dfs2(root, depth, L, parent):
ans_list[root] = ans_list[L] - depth
ans_node_list[root] = ans_node_list[L]
for e in adj_list[root]:
if e != parent:
dfs2(e, depth+1, L, root)

t = int(input())

for _ in range(t):
n, k, a = map(int,input().strip().split())
a -= 1
f = map(int,input().strip().split())
sp_nodes = [False]*n
for i in f:
sp_nodes[i-1] = True

adj_list = [[] for i in range(n)]
parent_list = [-1]*n
ans_list = [-1] * n
ans_node_list = [-1] * n

for _ in range(n-1):
    u, v = map(int,input().strip().split())
    u -= 1
    v -= 1
    adj_list[u].append(v)
    adj_list[v].append(u)


dfs(a, 0)

for i in range(n):
    if ans_list[i] == -1 and parent_list[i] != -1 and ans_node_list[parent_list[i]] != -1:
        dfs2(i, 1, parent_list[i], parent_list[i])


for i in range(n):
    ans_node_list[i] += 1
print(*ans_list)
print(*ans_node_list)

https://www.codechef.com/viewsolution/45790437
I did not change any logic just took the whole input at once and printed the whole output at once and it got AC. It is sometimes the case with python . I also learned this trick by reading blogs about good pratices for python in CP.

1 Like

@shubham1010 Please Format Your Code So That Everybody Could Understand It
Surround Your Code With ```

from sys import setrecursionlimit
setrecursionlimit(10**8)

def dfs(root, depth, parent = -1):
sp = -1
parent_list[root] = parent
for e in adj_list[root]:
if e != parent_list[root]:
temp = dfs(e, depth + 1, root)
if temp != -1 and sp == -1:
sp = temp

if sp_nodes[root]:
    sp = root
if sp != -1:
    ans_list[root] = depth
    ans_node_list[root] = sp

return sp
def dfs2(root, depth, L, parent):
ans_list[root] = ans_list[L] - depth
ans_node_list[root] = ans_node_list[L]
for e in adj_list[root]:
if e != parent:
dfs2(e, depth+1, L, root)

t = int(input())

for _ in range(t):
n, k, a = map(int,input().strip().split())
a -= 1
f = map(int,input().strip().split())
sp_nodes = [False]*n
for i in f:
sp_nodes[i-1] = True

adj_list = [[] for i in range(n)]
parent_list = [-1]*n
ans_list = [-1] * n
ans_node_list = [-1] * n

for _ in range(n-1):
    u, v = map(int,input().strip().split())
    u -= 1
    v -= 1
    adj_list[u].append(v)
    adj_list[v].append(u)


dfs(a, 0)

for i in range(n):
    if ans_list[i] == -1 and parent_list[i] != -1 and ans_node_list[parent_list[i]] != -1:
        dfs2(i, 1, parent_list[i], parent_list[i])


for i in range(n):
    ans_node_list[i] += 1
print(*ans_list)
print(*ans_node_list)

Try Reducing Recursion, Loops

Thank you. Please share some blogs if you have saved them.