I am trying to solve 1029E-Tree with Small Distances problem using DP on Tree but I am getting WA. My idea is:

If I am at some vertex u with distance dis from 1 (root) and every child of u is v except the parent then I can do the following stuffs:

• if dis==1 then visit all v with dis==2

• if dis==2 then either I can visit v with dis==3 or I can connect this vertex with root and then visit v with dis==2

• if dis==3: -- if u is a leaf then it must be connected to root

  -- or I can just visit v with dis==4 or I can connect u to root and visit v with dis==2

• id dis==4 then I must connect u to the root and visit v with dis==2

Here is my code : Solution