PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy - Medium
PREREQUISITES:
Dynamic programming, combinatorics
PROBLEM:
You’re given an array A such that \sum A_i = N.
Count the number of labeled trees on N vertices such that the following condition holds:
- You can pick two vertices u, v such that dist(u, v) = 2, and perform A_u \gets A_u+1, A_v\gets A_v-1.
- It’s then possible to convert A to [1, 1, \ldots, 1].
EXPLANATION:
Naturally, the first step is to understand when a fixed tree allows for converting A to [1, 1, \ldots, 1].
Let’s root the tree at vertex 1.
Call vertex u even if dist(1, u) is even, and odd otherwise.
The parity of a vertex denotes whether it is even or odd.
Observe that our given operation only allows us to transfer value between two vertices of the same parity, since if we operate on u and v then either dist(1, u) = dist(1, v) or |dist(1, u) - dist(1, v)| = 2 must hold.
So, if we denote:
- E_c and O_c to be the count of even/odd vertices, respectively.
- E_s and O_s to be the sum of even/odd vertices, respectively.
Then because there’s no interaction between parities, we surely need E_c = E_s and O_c = O_s to hold.
(In fact, note that if E_c = E_s then we automatically get O_c = O_s because \sum A_i = N).
This condition is not just necessary, but is also sufficient - if E_c = E_s then we are able to freely rearrange values among even vertices using our operation as long as the overall sum is maintained; thus allowing us to reach [1, 1, \ldots, 1] among them (and then the same applies to odd vertices.)
This is because, if u and v are two even vertices, we can always transfer one unit of value from u to v by repeatedly operating along the path from u to v.
Now that we have a characterization, let’s try to count the number of valid trees.
We’ll do this in two stages: first, we’ll split vertices into even/odd groups, then we’ll figure out how to combine these two groups into trees.
First, consider fixing the subset of even vertices.
A subset S of \{1, 2, \ldots, N\} can be the set of even vertices if and only if:
- 1 \in S, and
- The size of S equals the sum of values of vertices in S.
The number of such subsets can be computed using dynamic programming, by simply tracking the size and sum so far.
That is, define dp(i, x, y) to be the number of subsets of \{1, 2, \ldots, i\} that:
- Contain vertex 1,
- Have a size of x, and
- Have a sum of values equal to y.
Transitions are simple: we either include element i (adding 1 to size and A_i to sum), or don’t include it (adding 0 to both) and hence
with the base state being dp(1, 1, A_1) = 1 since we are forced to take vertex 1.
Because \sum A_i = N, this DP runs in \mathcal{O}(N^3) time.
It can easily be implemented using \mathcal{O}(N^2) memory if needed since we only care about the dp(N, \cdot, \cdot) values.
Once we know subset counts, let’s try to count ways to connect them into a tree.
Suppose we fix M to be the size of the even vertex set.
There are dp(N, M, M) ways to choose such a subset.
Once the subset is chosen, we have M vertices on one side, (N-M) vertices on the other, and we want to figure out how to join them into a tree while ensuring that edges are only added between vertices of different sides.
One way to view this is: we have the complete bipartite graph K_{M, N-M}, and we want to count the number of spanning trees it has.
It turns out that the answer to this question is a rather simple expression, simply being
The proof of why can be found below.
Thus, after computing the DP, the solution is simply to output
Proof of the formula
Perhaps the “simplest” proof of this fact, without resorting to much heavy machinery, is to use Prüfer codes.
Recall that the Prüfer code corresponding to a tree is obtained by repeatedly removing the smallest-labeled leaf and writing down the label of its neighbor, till only two vertices remain.
Thus, the Prüfer code of a tree with N vertices is a sequence of length N-2 with elements in [1, N].It can be proved that this is a bijection, which then shows that the number of labeled trees on N vertices is exactly N^{N-2}.
We modify this argument slightly for our purposes.
Rather than a single sequence, we build two sequences P_E and P_O.
These are initially both empty. Then, repeat the following:
- Choose the leaf with smallest label.
- If it is even, its (unique) neighbor is odd; append this neighbor to P_O.
- Alternately, if it is odd, its (unique) neighbor is even; append this neighbor to P_E.
Stop the process when only two vertices remain.
Observe that if we start with N odd vertices and M even vertices, then P_O will have size N-1 and P_E will have size M-1 in the end.
We’ll call the pair (P_O, P_E) the bipartite Prüfer code of the tree.Observe that there are M^{N-1} \cdot N^{M-1} possible values that the bipartite Prüfer code can take.
It’s not hard to now adapt the argument for ordinary Prüfer codes to bipartite Prüfer codes, to show that there’s a bijection between trees and possible codes - thus making the number of possible trees exactly M^{N-1} \cdot N^{M-1} as claimed.
TIME COMPLEXITY:
\mathcal{O}(N^3) per testcase.
CODE:
Editorialist's code (C++)
// #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;
vector pw(505, vector(505, 0));
for (int n = 1; n < 505; ++n) {
pw[n][0] = 1;
for (int i = 1; i < 505; ++i) {
pw[n][i] = (1ll* pw[n][i-1] * n) % mod;
}
}
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector a(n, 0);
for (int &x : a) cin >> x;
vector dp(n+1, vector(n+1, 0));
dp[1][a[0]] = 1;
for (int i = 1; i < n; ++i) {
for (int sz = i+1; sz >= 1; --sz) {
for (int sm = a[i]+a[0]; sm <= n; ++sm) {
dp[sz][sm] = (dp[sz][sm] + dp[sz-1][sm-a[i]]) % mod;
}
}
}
int ans = 0;
for (int i = 1; i < n; ++i) {
int ways = dp[i][i];
int trees = (1ll * pw[i][n-i-1] * pw[n-i][i-1]) % mod;
ans = (ans + 1ll*ways*trees) % mod;
}
cout << ans << '\n';
}
}