Consider DP, and try to record whether each subtree can be completely deleted. When transferring, it is required that when deleting the edge between \(u\) and one of its children \(v\), the colors of \(u\) and \(v\) are different. Thus, we design the DP as follows: \(f_{u,i}=0/1\) indicates whether there exists a valid operation sequence that completely deletes the subtree of \(u\), and it is predetermined that when deleting the edge between \(u\) and its parent, the color of \(u\) is \(i\).

Specifically, we require that at a certain moment when the color of \(u\) is \(i\), we insert a step, which corresponds to deleting the edge \((u, f a_u)\).

Consider the transition, which involves assigning an order to all children such that when each edge is deleted, the colors at both ends of the edge are different.

We find that this is essentially a matching problem, as we can directly determine how many times the color of \(u\) is \(0, 1, 2\) respectively (i.e., the number of occurrences in the sequence \(c_u,(c_u+1) \bmod 3,\cdots,(c_u+\deg_u-1) \bmod 3\), where if we are computing \(f_{u,i}\), the count for \(i\) is reduced by \(1\)). This then becomes a bipartite graph multiple matching problem with \(3\) nodes on the left and \(\deg_u-1\) nodes on the right.

We can directly run a standard bipartite graph matching algorithm, with time complexity \(O(\sum \deg_u^2)=O(N^2)\). In practice, since the left side has only \(O(1)\) nodes, we can use Hall’s theorem to achieve a time complexity of \(O(N)\).

When constructing the solution, we can simply iterate over each right-side node and match it with a left-side node, as long as it remains valid.

https://atcoder.jp/contests/abc437/submissions/71954949 (I used map so this is \(O(n\log n)\), it is easy to do this in \(O(n)\))