Check if a tree is formed by 3 subtrees with given number of nodes

Question

I have run into a contest problem (ACM like) that sounds like this:

Input: a tree of $N$ nodes; integers $X,Y,Z$ such that $X+Y+Z=N$
Question: Can the tree be partitioned into three trees of $X,Y,Z$ nodes by removing only two edges?

I am promised that $N \le 10^5$. My code must finish within 200ms.

The brute-force method of generating all possible partitions of X, all possible partitions of Y, and all possible partitions of Z and checking if they form the tree is too slow. This got me thinking about dynamic programming. The only thing I see here to help me form a DP algorithm is counting how many children does a node have (treating the current node as a root of a tree). However, checking between the nodes themselves for X, Y, Z would still take $O(N^3)$ time, which is too much.

So do you see any other algorithm for this task?

Example: We're given a tree of 8 nodes and X = 3, Y = 3, Z = 2 and the following edges: 2-3; 4-2; 2-1; 1-5; 1-8; 5-6; 5-7. The answer is yes because you can delete edges 1-2 and 1-5 to form 3 partitions like this {1,8} (Z) and {2,3,4}/{5,6,7} (X/Y).

This problem is from a preparation stage of a local competition that happend a while back.

Answer 1

Let's root the tree arbitrarily at node $r$. Now with a DFS, compute the size of each subtree rooted at each node $v$. Let $good(u,i)$ be true, if there is some subtree rooted at $v$ where $v$ lies in the subtree rooted at node $u$ such that the size of the subtree rooted at $v$ is exactly $i$. We will maintain this property for only three values of $i \in \{X,Y,Z\}$. This can be done in the same DFS itself (which is left as an exercise for you).

Since we have rooted the node at $r$, due to the DFS traversal from $r$, all nodes except $r$ have a parent which we will denote by $p(u)$ for a node $u$. Now let's compute the same function $good'(u,i)$ for the same 3 values of $i$ but instead for the subtree which is obtained after deleting the subtree rooted at node $u$. This value can be computed via a second DFS (which I also leave as an exercise for you).

HINT: Read up on how to find the longest and second longest paths starting at some node in a unrooted tree for each node.

Now the answer to the question is "Yes" if $\exists\ u \in$ tree $T$ such that $good(u,i)$ is true and $good'(u,j)$ is true and $i \not= j$

Since we have only two DFS traversals, the time complexity is $\mathcal{O}(n)$ and a constant amount of memory per node implies $\mathcal{O}(n)$ memory as well.