PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
DFS
PROBLEM:
You’re given a tree.
A subset S of vertices is called good if there exists some vertex x \in [1, N] such that for each y, z \in S with y \ne z, x lies on the path from y to z.
Note that x need not belong to S.
Count the number of good vertex subsets in the tree.
EXPLANATION:
Let’s understand when a subset of vertices S can be good.
If S has size 1 or 2, then it is clearly good: choosing any element of S will be a valid choice of x, for example.
Let’s now focus on |S| \ge 3.
The main observation that allows us to solve the problem now is that for such a subset, there can be at most one choice of x that allows for it to be good.
To see why: let’s pick three distinct vertices v_1, v_2, v_3 \in S.
Root the tree at v_1, and consider the paths v_1 \to v_2 and v_1 \to v_3.
x must lie on both of these paths.
So, if L = \text{lca}(v_2, v_3) when rooted at v_1, the only possible choices for x are L and its ancestors.
On the other hand, x must also lie on the v_2 \to v_3 path.
This means the only possible choice for x is L itself!
Of course, since S may have other vertices, it’s possible that this x ends up being invalid for some other pair; but either way we see that there can be at most one choice.
This allows us to count good subsets by characterizing them by their choice of x.
As noted earlier, any subset with size \le 2 is good.
There are N + \binom{N}{2} such subsets; so we add this count to the answer.
Next, let’s fix a vertex x, and try to count the number of subsets of size \ge 3 such that all paths between pairs of vertices pass through x.
For a path between u and v to pass through x, we must have the following condition:
- Delete all edges connected to x from the tree.
This will create several connected components - specifically, there will be \text{deg}(x) + 1 components.
Note that x itself will be an isolated vertex; we aren’t deleting it, only its edges. - Then, if u and v belong to different connected components, their path in the tree passes through x; otherwise it does not.
Now, for a general subset S, this tells us that there can be at most one vertex present in each component after deleting edges connected to x.
So, upon fixing x, our aim is to count the number of ways of picking at least three vertices, while also ensuring we pick at most one from each component.
Suppose there are k = \text{deg}(x) + 1 components, with sizes s_1, s_2, \ldots, s_k.
Let’s define a function \text{ways}(i, j) where 1 \le i \le k and 0 \le j \le 3 as follows:
- \text{ways}(i, j) for j \le 2 denotes the number of ways of picking exactly j vertices from the first i components; at most one vertex per component.
- \text{ways}(i, 3) denotes the number of ways of picking at least 3 vertices from the first i components; at most one vertex per component.
With this definition, computing these values is fairly easy:
- \text{ways}(i, 0) = 1 always, there’s only one way to pick an empty subset.
- \text{ways}(i, 1) = \text{ways}(i-1, 1) + s_i
We either pick one vertex from the i-th component (s_i choices) or from one of the previous i-1 components. - \text{ways}(i, 2) = \text{ways}(i-1, 2) + s_i \cdot \text{ways}(i-1, 1)
We either pick one vertex from the i-th component and one from among the previous i-1, or ignore this component and pick two from the previous i-1. - \text{ways}(i, 3) = \text{ways}(i-1, 3) + \text{ways}(i-1, 3) \cdot s_i + \text{ways}(i-1, 2) \cdot s_i
We either pick \ge 3 vertices from the previous i-1 components and nothing from the i-th, or \ge 3 vertices previously and any one vertex from the i-th, or two vertices previously and one from here.
This allows us to compute all \text{ways}(i, j) values in \mathcal{O}(k) time.
Once this is done, the number of valid subsets is exactly \text{ways}(k, 3) by definition.
Note that since k = \text{deg}(x) + 1, when summed up across all x the summation becomes
This is because \sum \text{deg}(x) counts each edge exactly twice (once at each endpoint) and there are N-1 edges.
Thus, the complexity is \mathcal{O}(N).
The only remaining part is for us to actually get the component sizes s_1, s_2, \ldots, s_k for a fixed vertex x.
However, this is not too hard.
Let’s root the tree at 1, and define \text{subsz}[u] to be the size of the subtree of u.
Then, for a vertex x whose children are c_1, c_2, \ldots, c_m,
- For each i = 1, 2, \ldots, m, the child subtree of c_i will be one component.
This has a size of \text{subsz}[c_i] - The vertex x itself will be a single component, with size 1.
- The entire rest of the tree will be a single component, having size N-\text{subsz}[x].
Knowing the sizes, we can thus run the previous algorithm to compute the answer.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
const int mod = 998244353;
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector adj(n, vector<int>());
for (int i = 0; i < n-1; ++i) {
int u, v; cin >> u >> v;
adj[--u].push_back(--v);
adj[v].push_back(u);
}
vector subsz(n, 0);
int ans = (n + 1ll*n*(n-1)/2) % mod;
auto dfs = [&] (const auto &self, int u, int p) -> void {
subsz[u] = 1;
array<int, 4> ways = {1, 1, 0, 0};
auto upd = [&] (int x) {
ways[3] = (1ll * ways[3] * (x + 1) + 1ll * ways[2] * x) % mod;
ways[2] = (ways[2] + 1ll * x * ways[1]) % mod;
ways[1] = (ways[1] + 1ll * x * ways[0]) % mod;
};
for (int v : adj[u]) if (v != p) {
self(self, v, u);
subsz[u] += subsz[v];
upd(subsz[v]);
}
upd(n - subsz[u]);
ans = (ans + ways[3]) % mod;
};
dfs(dfs, 0, 0);
cout << ans << '\n';
}
}