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.
TL;DR: It might be a high constant factor.
Sets have a higher constant factor than many other data structures/ algorithms with $$$O(N log N)$$$ time complexity (I've heard that the reason is because of all the pointers used, but I don't know how true that is).
Also, I'm not too familiar with Codechef, but it seems like it's only TLE'ing for one test case, by a slim margin.