This is in relation with the problem Special Tree from April Lunchtime 2021. Problem Link.
I used multisource BFS/Dijkstra to solve the problem in O(NlogN). Here are the two submissions I made.
The only difference between TLE solution and Accepted Solution is that, in TLE solution I added two extra sets (declared at line 96 and 97, updated at line 102 and 106), which I thought would be useful ahead in the program, but were never used. Whereas in the Accepted solution I removed those two sets. This is the only difference.
The max length of those two sets is N and it would create a difference of NlogN only, which is well under limits.
Is adding elements to a set that much costly ? I think I am missing out something here.