PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: satyam_343
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
DFS, Dynamic programming
PROBLEM:
You’re given a tree.
For a subset S of vertices, define f(S) to be the LCA-closure of S.
For each K = 1, 2, \ldots, N, compute the number of subsets for which |f(S)| = K.
EXPLANATION:
Let’s understand which vertices will be present in f(S) for a fixed S.
Consider some vertex u.
If u \in S, then certainly u \in f(S).
If u \not \in S, then u \in f(S) if and only if there are two vertices v and w such that v, w \in S and \text{lca}(v, w) = u.
(Note that we don’t even have to split into these two cases really; since the second case includes the first by choosing v = w = u).
In fact, the set f(S) is exactly the set of vertices that appear in the virtual tree of S.
While this knowledge is not really relevant to solving this problems, if you’re interested in learning about the concept, here is a tutorial.
We can use the fact that u being present depends only on the state of vertices in its subtree, to write a dynamic programming solution.
Define dp[u][x] to be the number of subsets of vertices in the subtree of u, whose virtual tree has size x.
There are two possibilities:
- All vertices in the subtree of u, lie entirely in the subtree of one of its children.
In this case, u is not forced to be in f(S), and so it will be present only if it is itself chosen to be in S. - The subset has vertices lying in two different child subtrees of u.
In this case, u is forced to be in f(S).
The first case is simple to deal with.
For each child c,
- Add dp[c][x] to dp[u][x], to account for not taking u into S.
- Add dp[c][x-1] to dp[u][x], to account for taking u into S.
Be careful about the contribution of dp[c][0] though, since the empty set and \{u\} must be counted only once each.
Next, we need to deal with the second case, i.e. choosing vertices from different children (which forces u to be present).
Let c_1, c_2, \ldots, c_k be the children of u, and suppose we have subsets of size x_1, x_2, \ldots, x_k from each of them.
This results in a set with overall size (x_1 + x_2 + \ldots + x_k) + 1, since u is forcibly included in f(S).
However, we need to ensure that at least two of the x_i values are \ge 1, since otherwise everything will lie in a single subtree (and hence u is not forced).
Moreover, note that each such subset must be counted twice, since we can freely choose to include/not include u itself in S (u will always be in f(S), but we’re counting S and not f(S).)
This can be handled as follows.
Let’s define a temporary dp, temp[u][x].
Initially, temp[u][0] = 1 and everything else is 0.
Now, for each i = 1, 2, \ldots, k, merge the subtree of c_i into temp.
That is, for each x and y, update temp[u][x+y] \gets temp[u][x+y] + temp[u][x]\cdot dp[c_i][y]
In the end, temp[u][x] will hold the number of ways of choosing a subset for which |f(S)| = x from the subtree of u (without accounting for u).
From this, we can then subtract out the “bad” subsets, i.e. those where all elements lie in a single child.
This is simple enough: just subtract dp[c_i][x] from temp[u][x] for each child c_i and each value of x.
Finally, we can update dp[u][x], by incrementing dp[u][x+1] by 2\cdot temp[u][x].
The multiplicative factor is for deciding whether to include u or not, and moving x to x+1 is to signify the inclusion of u in f(S).
We have \mathcal{O}(N^2) states; however, the transitions appear slow at first glance.
Luckily however, they are not.
Note that in our transitions, the dimension we iterate over does not exceed the size of the subtree under consideration.
In such cases, it can be proved that the overall work done, across the entire tree, remains quadratic.
To understand why, see point 7 of this blogpost.
So, our overall solution is quadratic which is fast enough to pass.
TIME COMPLEXITY:
\mathcal{O}(N^2) 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());
const int mod = 998244353;
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
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);
}
auto dfs = [&] (const auto &self, int u, int p) -> vector<int> {
vector<int> res = {1}, sub = {1};
for (int v : adj[u]) if (v != p) {
auto ch = self(self, v, u);
for (int x = size(res)-1; x >= 0; --x) {
for (int y = 1; y < size(ch); ++y) {
if (x+y == size(res)) res.push_back(0);
res[x+y] = (res[x+y] + 1ll * res[x] * ch[y]) % mod;
}
}
for (int i = 1; i < size(ch); ++i) {
if (i == size(sub)) sub.push_back(0);
sub[i] = (sub[i]+ ch[i]) % mod;
}
}
for (int i = 0; i < size(res); ++i) {
if (i < size(sub)) res[i] = 2ll*(res[i] + mod - sub[i]) % mod;
else res[i] = 2*res[i] % mod;
}
res.insert(res.begin(), 0);
for (int i = 0; i < size(res); ++i) {
if (i < size(sub)) res[i] = (res[i]+ sub[i]) % mod;
if (i > 1 and i <= size(sub)) res[i] = (res[i] + sub[i-1]) % mod;
}
if (size(res) == 1) res.push_back(1);
else res[1] += 1;
return res;
};
auto dp = dfs(dfs, 0, 0);
for (int i = 1; i <= n; ++i) cout << dp[i] << ' ';
cout << '\n';
}
}