The three possible values of ~K~ are very different problems. We will present their solutions in the order ~K = 1, K = 3, K = 2~ as this is roughly the increasing order of difficulty.

Subtask 1: ~K = 1~

When ~K = 1~, the answer is the longest node-weighted path from the root, which can be found in a single DFS.

Subtasks 5, 6: ~K = 3~

When ~K = 3~, it turns out that it is always possible to visit all nodes. We can do this recursively using the state ~\text{DFS}(u)~, which starts from a node ~u~, traverses every node in the subtree of ~u~ once, and ends at a child of ~u~ (or ~u~ itself if the subtree consists of just ~u~). If ~v_1,\dots,v_k~ are the children of ~u~, ~\text{DFS}(u)~ represents the path formed by ~u \to \text{rev}(\text{DFS}(v_1)) \to \text{rev}(\text{DFS}(v_2)) \to \dots \to \text{rev}(\text{DFS}(v_k))~, where ~\text{rev}~ denotes reversing the path.

Reconstructing the path naively (copying around and reversing parts of the path) takes ~\mathcal{O}(N^2)~ time and only passes subtask 5. It is possible to use small to large merging of double-ended queues to solve subtask 6 in ~\mathcal{O}(N \log N)~ time. We can also reconstruct the path in linear time by instead building up a second graph structure representing parts of the path, only passing around the endpoints of the partial path during the reconstruction DFS, and doing a path traversal afterwards to print the path in the correct order.

Note that this solution would also solve all of ~3 \le K \le N~ if such ~K~ were permitted.

Subtasks 2, 3, 4: ~K = 2~

When ~K = 2~, we need to use dynamic programming to find the optimal path. The following solution may be found by very carefully analyzing all possible ways the path may go. After writing some dynamic programming code to calculate the answer, it is recommended to write a brute force solution and compare answers with it on many random tests before trying to write the path reconstruction code, as it is very easy to miss a case in this problem.

Let ~dp1[u]~ represent the best path starting at ~u~ and staying within the subtree of ~u~. Let ~dp2[u]~ represent the best path starting at ~u~, staying within the subtree of ~u~, and ending at a child of ~u~ or just ~u~. Let ~dp3[u]~ represent the best path starting at a child of ~u~ and staying within the subtree of ~u~ (or just ~u~). Let ~v_1,\dots,v_k~ be the children of ~u~. The transitions are of the form:

• ~dp1[u] \leftarrow p[u] + dp2[v_1] + p[v_2] + \dots + p[v_{k-1}] + dp1[v_k]~, representing the path ~u \to \text{rev}(dp2[v_1]) \to v_2 \to \dots \to v_{k-1} \to dp1[v_k]~
• ~dp1[u] \leftarrow p[u] + dp3[v_1]~, representing the path ~u \to dp3[v_1]~
• ~dp2[u] \leftarrow p[u] + dp2[v_1] + p[v_2] + \dots + p[v_{k-1}] + p[v_k]~, representing the path ~u \to \text{rev}(dp2[v_1]) \to v_2 \to \dots \to v_{k-1} \to v_k~
• ~dp3[u] \leftarrow p[v_1] + p[v_2] + \dots + dp2[v_{k-1}] + p[u] + dp3[v_k]~, representing the path ~v_1 \to v_2 \to \dots \to dp2[v_{k-1}] \to u \to dp3[v_k]~
• ~dp3[u] \leftarrow dp2[v_1] + p[u] + dp2[v_2] + p[v_3] + \dots + p[v_{k-1}] + dp1[v_k]~, representing the path ~dp2[v_1] \to u \to \text{rev}(dp2[v_2]) \to v_3 \to \dots \to v_{k-1} \to dp1[v_k]~

and similar ones for other orderings of the children.

These transitions mostly just add up ~p[v_i]~, except up to ~3~ of the children are replaced with a dp value. For subtask 2, it suffices to loop over all ~\mathcal{O}(k^3)~ ways to choose these up to ~3~ special children for a time complexity of ~\mathcal{O}(N^3)~. For subtasks 3 and 4, we can instead maintain the ~3~ children ~v~ with the maximum ~dp[v]-p[v]~ and only consider these to compute the dp values in ~\mathcal{O}(N)~ time. Reconstructing the path in ~\mathcal{O}(N^2)~ time naively suffices for subtask 3, and for subtask 4, an ~\mathcal{O}(N \log N)~ or ~\mathcal{O}(N)~ time reconstruction is needed, which can be done similar to the ~K = 3~ case.