Virtual Tree

Often, when running operations or dynamic programming on trees, we only need to keep track of a few key nodes and their relationships. If we are able to isolate the few relevant nodes that keep the structure of the tree intact, while still maintaining the relationships between key nodes, we can cut our time complexity down significantly.

Example - Leaf Color

Naive Solution

We can sum the number of vertex sets $T$ that satisfy the condition with degree $1$ vertices having color $c$ across all $c \in [1,N]$ to obtain an answer.

We define $\texttt{dp}[v]$ to be the number of induced subgraphs that are trees rooted at $v$ such that all leaves have color $c$. When transitioning, initialize $\texttt{dp}[v] = 1$. Transitions are as follows:

• For every child $u$ of $v$, we can either attach one of the induced subgraphs that are trees containing $u$ or ignore $u$ completely:

  $$\texttt{dp}[v] = \texttt{dp}[v] \cdot (1 + \texttt{dp}[u])$$

• However, if $a_v \neq c$, then we have to remove the case where the induced subgraph is just $v$ itself, because this violates the condition:

  $$\texttt{dp}[v] = \texttt{dp}[v] - 1$$

To extract our answer, we can iterate over every possible root $r$ of our induced subgraph.

• If $a_r =c$, then we can directly add $\texttt{dp}[v]$ to our answer.
• Otherwise, we need to make sure that $r$ is not degree $1$ in any of the induced subgraphs we add. Thus, we subtract all the cases when $r$ is attached to exactly one of its children.

Thus, our answer will be

The time complexity for this approach is $\mathcal{O}(n^2)$, because for each color, we iterate over all $n$ nodes.

Here is a (TLE) code that demonstrates this approach:

Copy
#include <bits/stdc++.h>
using namespace std;

using ll = long long;

constexpr int MAX_N = 2e5 + 1;
constexpr int MOD = 998244353;

int n, a[MAX_N];
ll dp[MAX_N], ans;

Virtual Tree Optimization

To accelerate our dynamic programming approach, we can observe that our dynamic programming transitions are relatively simple when involving certain types of nodes.

The most simple case of this is when a node $v$ has no nodes of color $c$ in its subtree. Here, we can simply conclude $\texttt{dp[v]}=0$, because it is impossible to construct an induced subgraph in the subtree of $v$ with leaves of color $c$.

Another case of simple transitions is when only one child $u$ of $v$ has any nodes of color $c$. Then, from our previously established dp transitions, we observe that $\texttt{dp}[v]$ is equivalent to $\texttt{dp}[u]$.

Let $S$ be the set of nodes excluding these two cases. More formally, if we let $S_c = \{ v \in V \mid a_v = c \}$, then

We define $S$ to be the virtual tree or auxillary tree of $S_c$. We can also define $p'_v$, the virtual parent of $v$ as the lowest node in $S$ such that $p'_v \neq v$ and $\text{LCA}(p'_v,v) = p'_v$, and $v$ as a virtual child of $p'_v$.

Here is an example of a virtual tree, with nodes in $S_c$ colored purple, and the rest of the virtual tree in blue:

From here, we can calculate our dp values similarly to before, except we only process nodes in $S$, and we consider virtual children instead of direct children.

Notice that when calculating our answer, we also only need to consider nodes in $S$. Nodes not in $S$ can not be the root because only one of their children have nodes of color $c$ in their subtree, meaning that we will end up with a degree $1$ node not of color $c$.

Thus, our time complexity will be $\mathcal{O}(\sum_c |S|)$, since we iterate over the set $S$ for each distinct color.

In the following section, we prove that $|S|$ is bounded by $\mathcal{O}(S_c)$, and also demonstrate how to construct a virtual tree.

Virtual Tree Construction

A well-known method to construct the virtual tree of set $S_0$ is as follows:

• Sort the points in $S_0$ by DFS order, and add them to $S$
• Calculate the LCA of any two adjacent key points in $S_0$, and add that to $S$
• Build the virtual tree based on the ancestor-descendant relationship of the original tree

For a rough proof of correctness, consider the following:

Observe that a vertex $v \in S$ must either satisfy $v \in S_0$ or there exist at least two distinct children of $v$ with key points in their subtrees.

The former is taken care of immediately in step 1. In the latter case, as a property of DFS order, there must exist key points $a,b$ in the subtree of $v$ such that $a,b$ are in distinct subtrees and $a,b$ are adjacent in the DFS order when only considering key points. Because $v$ is $\text{LCA}(a,b)$, our construction neccesarily holds true. A corollary of our construction is that $|S|$ is indeed bounded by $\mathcal{O}(|S_0|)$ - in fact, a tight bound is $|S| \leq 2 \cdot |S_0|-1$.

Implementation

Below is an implementation for the task, which also contains general code for constructing virtual trees. Because $\sum_c |S|$ is $\mathcal{O(n)}$, and we sort by DFS order, our final time complexity is $\mathcal{O}(n\log n)$.

Copy
#include <bits/stdc++.h>
using namespace std;

using ll = long long;

constexpr int MAX_N = 2e5 + 1;
constexpr int LG = 18;
constexpr int MOD = 998244353;

int n, a[MAX_N];

Problems