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

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.

• Hint: Start by trying to solve this problem for the case of two trees and removing only one edge. What can you come up with? – D.W. Mar 1 '16 at 18:38

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.

• I think this is close enough as a hint to the question. There may be something missing: What if one subtree of desired size is in the left subtree with another size found in the right? What if one good tree is part of another? (A peeve of mine: Don't call every traversal <Whatever> First Search. For summarising subtree information, post-order seems fit.) (Oh, and one of the final trees will be rooted at the original root and not be a subtree of the original one - if you find a more accurate title that also makes a better title…) – greybeard Mar 3 '16 at 6:11
• Actually all those special cases can be avoided if you keep track of what properties you get from which child for a particular node, like you do to find the longest and second longest paths starting at each node and then carefully combining them in the final step :) – Banach Tarski Mar 3 '16 at 6:33