title: Virtual Trees

tags:

- Original

This article explores the concept of *virtual trees*, proves their key properties, describes an efficient algorithm for their construction, and then applies the method to solve a specific problem using dynamic programming (DP) on trees.

Many tree-related problems require working with a subset of vertices while preserving the tree structure induced by their pairwise lowest common ancestors (LCA). The concept of a *virtual tree* allows us to transition from the original tree $T$ with $N$ vertices to a substructure whose size linearly depends on the size of the selected set $X$. This significantly accelerates algorithms, particularly for DP computations.

Let $T$ be a given rooted tree, and $X$ be some subset of its vertices. A **virtual tree** $A(X)$ is defined as follows:

where $\operatorname{lca}(x,y)$ denotes the lowest common ancestor of vertices $x$ and $y$ in tree $T$. In this tree, an edge is drawn between each pair of vertices if one of them is an ancestor of the other in $T$.

This definition ensures that the structure includes all vertices "important" for analyzing $X$ (i.e., all LCAs for pairs of vertices from $X$), and connectivity is inherited from the original tree.

**Statement.**

For any set $X$ of vertices of tree $T$, the following inequality holds:

**Proof.**

Let's choose a depth-first search (DFS) traversal order and consider the vertices of set $X$ arranged in traversal order:

Denote

We need to prove the following lemma.

!!! note "Lemma"

For any integer $n \ge 3$, if $\operatorname{lca}(x_1,x_n) \neq l$, then $\operatorname{lca}(x_1,x_n) = \operatorname{lca}(x_2,x_n).$

??? hint "Proof"

Suppose that for some $n \ge 3$, both $\operatorname{lca}(x_1,x_n) \neq l$ and $\operatorname{lca}(x_1,x_n) \neq \operatorname{lca}(x_2,x_n)$ are true. Let $k = \operatorname{lca}(x_1,x_n)$.

If $k$ is an ancestor of $l$, then by definition both $\operatorname{lca}(x_1,x_n)$ and $\operatorname{lca}(x_2,x_n)$ equal $k$, which contradicts our assumption. Therefore, $k$ must be a descendant of $l$. But then in the DFS order, we would get the sequence $x_1, x_n, x_2$, which contradicts our chosen traversal order.

This contradiction completes the proof of the lemma.

Returning to the proof of the property, we note that according to the lemma, each LCA arising from sequential consideration of the vertices in $X$ equals either:

1. the vertex $x_1$ itself,

2. $\operatorname{lca}(x_1,x_2)$, or

3. $\operatorname{lca}(x_2,x_n)$ for some $n$.

Thus, we can recursively write:

From this, it follows that

Applying induction on the size of set $X$, we obtain the final estimate:

Moreover, if tree $T$ is a perfect binary tree and $X$ consists of its leaves, the inequality is achieved exactly.

**Statement.**

Let the vertices of $X$ be ordered by DFS traversal order: $x_1, x_2, \dots, x_m$. Then

**Proof.**

We start with the recursive representation:

Let's expand it recursively:

Thus, we obtain the required representation.

Given a set of vertices $X$, the virtual tree $A(X)$ can be constructed in time

with preprocessing of the original tree $T$ in $O(N \log N)$ to ensure fast LCA queries.

**Main construction steps:**

1. **Preprocessing.**

Using depth-first search (DFS) or methods such as binary lifting or HLD, we compute entry times for vertices and LCAs for any two vertices.

2. **Vertex Sorting.**

We sort the vertices of set $X$ in the order of their appearance during DFS traversal (using entry times).

3. **Adding LCAs.**

For consecutive vertices $x_i$ and $x_{i+1}$, we compute $\operatorname{lca}(x_i,x_{i+1})$. The union of $X$ and these LCAs gives the set of vertices $A(X)$ (according to Property 2).

4. **Tree Construction.**

Having obtained the set of vertices $A(X)$ sorted by DFS order, we can traverse the sequence using a stack to restore the tree structure. During traversal, we maintain a stack whose top element is the current vertex, and if a new vertex is not a descendant of the vertex at the top of the stack, elements are "popped" until the corresponding ancestor is found, after which connecting occurs.

This algorithm guarantees that the virtual tree contains $O(|X|)$ vertices, allowing efficient application of DP.

Given a tree $T$ with $N$ vertices numbered from $1$ to $N$. Edge $i$ connects vertices $u[i]$ and $v[i]$. Each vertex $i$ is colored with color $A[i]$.

It is necessary to find (modulo 998244353) the number of (non-empty) subsets $S$ of vertices of tree $T$ satisfying the condition:

- The induced graph $G[S]$ is a tree.

- All vertices of $G[S]$ with degree 1 have the same color.

The main idea is to break down the problem by colors. For each color $c$, we consider the set of vertices $X$ colored with $c$ and build a virtual tree for $X$. Then, on the resulting tree, we perform DP to count valid subtrees where all vertices with degree 1 have color $c$. The final answer is obtained by summing the results for all colors.

Thanks to the construction of the virtual tree, although the original tree contains $N$ vertices, each virtual tree for a specific color has size $O(|X|)$, which allows DP to be performed in an acceptable time.

We get that the total complexity of preprocessing and building virtual trees does not exceed $O(N \log N + \sum_{c \in C} (|X_c| (\log |X_c| + \log N)) = O(N \log N)$, where $C$ is the set of different vertex colors and $X_c$ is the set of vertices with color $c$.

Below is the implemenation with comments describing the main components of the algorithm:

```{.cpp file=virtual_trees}

const int MOD = 998244353;

vector<int> g[MAXN]; // Adjacency list of the given graph

vector<int> vertex_sets[MAXN]; // For each color c, vertex_sets[c] stores vertices with color c

int tmr, n; // Global time counter and number of vertices

int up[LOGN][MAXN], dep[MAXN], tin[MAXN]; // For computing LCA and entry time

int col[MAXN]; // Vertex colors array

vector<int> virtual_g[MAXN]; // Adjacency list for the virtual tree

int dp[MAXN][2], sum[MAXN];

void dfs_precalc(int v, int p) {

tin[v] = ++tmr;

up[0][v] = p;

dep[v] = dep[p] + 1;

for (int i = 1; i < LOGN; ++i)

up[i][v] = up[i - 1][up[i - 1][v]];

for (auto to : g[v])

if (to != p)

dfs_precalc(to, v);

}

// Function to compute LCA using binary lifting

int getlca(int x, int y) {

if (dep[x] < dep[y]) swap(x, y);

for (int i = LOGN - 1; i >= 0; --i)

if (dep[up[i][x]] >= dep[y])

x = up[i][x];

if (x == y) return x;

for (int i = LOGN - 1; i >= 0; --i)

if (up[i][x] != up[i][y]) {

x = up[i][x];

y = up[i][y];

}

return up[0][x];

}

// DFS on the virtual tree to perform DP.

// Parameter c — target color for which counting is performed.

void dfs_calc(int v, int p, int c, int &ans) {

dp[v][0] = dp[v][1] = 0;

sum[v] = 0;

for(auto to : virtual_g[v]) {

if(to == p) continue;

dfs_calc(to, v, c, ans);

// DP transitions: combining current state with result from subtree.

int nxt0 = (dp[v][0] + sum[to]) % MOD;

int nxt1 = ((dp[v][0] + dp[v][1]) * 1ll * sum[to] % MOD + dp[v][1]) % MOD;

dp[v][0] = nxt0;

dp[v][1] = nxt1;

}

sum[v] = (dp[v][0] + dp[v][1]) % MOD;

if(col[v] == c) {

// If the vertex has the target color, it can participate in a valid subtree.

sum[v] = (sum[v] + 1) % MOD;

ans = (ans + sum[v]) % MOD;

} else {

ans = (ans + dp[v][1]) % MOD;

}

}

// Function to build a virtual tree for color c and perform DP.

void calc_virtual(int c, int &ans) {

auto p = vertex_sets[c];

if (p.empty()) return;

// Sort vertices by entry time (tin) — inorder traversal order.

sort(p.begin(), p.end(), [&](const int a, const int b) { return tin[a] < tin[b]; });

vector<int> stack = {1}; // Initialize stack with the root of tree T (vertex 1).

virtual_g[1].clear();

auto add = [&](int u, int v) {

virtual_g[u].push_back(v);

virtual_g[v].push_back(u);

};

// Process each vertex from set p, maintaining a stack to build the virtual tree.

for (auto u : p) {

if (u == 1) continue;

int lca = getlca(u, stack.back());

if (lca != stack.back()) {

while (stack.size() >= 2 && tin[lca] < tin[stack[stack.size() - 2]]) {

add(stack.back(), stack[stack.size() - 2]);

stack.pop_back();

}

if (stack.size() >= 2 && tin[lca] != tin[stack[stack.size() - 2]]) {

virtual_g[lca].clear();

add(stack.back(), lca);

stack.back() = lca;

} else {

add(stack.back(), lca);

stack.pop_back();

}

}

virtual_g[u].clear();

stack.push_back(u);

}

while (stack.size() > 1) {

add(stack.back(), stack[stack.size() - 2]);

stack.pop_back();

}

// Perform DP on the virtual tree, starting from root 1.

return dfs_calc(1, 0, c, ans);

}

// The main function where we read input data and calculate total answer

int solve(int N, const vector<pair<int, int>> &edges, const vector<int> &colors) {

n = N;

for(auto [x, y] : edges) g[x].push_back(y), g[y].push_back(x);

copy(colors.begin(), colors.end(), col + 1);

// Group vertices by color.

for(int i = 1; i <= n; ++i) vertex_sets[col[i]].push_back(i);

dfs_precalc(1, 0);

int ans = 0;

// Process the corresponding virtual tree for each possible color.

for (int i = 1; i <= n; ++i) calc_virtual(i, ans);

return ans;

}

- [Virtual Trees](graph/virtual_trees.md)

#include <cassert>

#include <algorithm>

#include <vector>

using namespace std;

namespace Implementation {

const int MAXN = 30;

const int LOGN = 10;

#include "virtual_trees.h"

}

int main() {

int n = 20;

vector<pair<int,int>> edges = {

{1,2}, {1,4}, {1,5}, {1,6}, {1,10},

{1,16}, {2,3}, {2,15}, {2,19},

{3,11}, {4,7}, {4,9}, {6,12},

{6,13}, {6,14}, {6,20}, {7,8},

{8,17}, {13,18}

};

vector<int> colors = {2, 10, 6, 3, 16, 20, 6, 17, 8, 13, 9, 11, 7, 12, 5, 13, 7, 18, 3, 18};

assert(Implementation::solve(n, edges, colors) == 25);

return 0;

}