PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming
PROBLEM:
You’re given a tree on N vertices.
Find the number of arrays A such that:
- A is a permutation of [0, N-1]
- For each non-leaf vertex u, A_u = \text{MEX}\{A_x \mid x \in S_u \setminus \{u\}\} where S_u is the subtree of u.
EXPLANATION:
Let’s try to just assign some values and see what happens.
Suppose we set A_1 = X, i.e. give the root the value X.
Note that no matter how we assign values to the other vertices, the MEX condition will be satisfied at the root - because A is a permutation, the values 0, 1, 2, \ldots, X-1 must appear at other vertices, while X won’t (since it’s at the root); so the MEX will certainly equal X.
Let’s next look at the children of the root.
Note that any children that are leaves can be functionally ignored - they can receive any value, so we can just fill them in at the end.
This means we only need to care about the children of the root that are not leaves.
There are now a few things that can happen:
- It’s possible that all children of the root are leaves, or in other words the tree is a star centered at 1.
In this case the answer is obviously N! since any arrangement will work - we already established that the condition is always true for the root. - Next, it’s possible that the root has exactly one child that’s not a leaf - let this be v.
The problem can then be solved recursively for v - though we have to keep in mind that the value X cannot be chosen, so v cannot receive a value that’s \gt X at all (since this would require X to be in the subtree of v, which is impossible). - Finally, it’s possible that the root has \ge 2 children that are not leaves.
In this case, it turns out that no valid configuration can exist!
To see why: suppose u and v are two non-leaf children.
Either A_u or A_v must be non-zero, since A is a permutation. Suppose A_u \gt 0.
This means that the value 0 appears in the subtree of u.
But then this also means A_v \ne 0, which in turn means v appears in the subtree of v - a contradiction since there can be only one 0, and the subtrees of u and v are disjoint.
It’s not hard to see that the last case above applies to any arbitrary vertex too, not just the root - that is, if any vertex has \ge 2 children that are non-leaves, then no valid configuration can exist, and the answer is immediately 0.
The above conclusion means we only need to concern ourselves with trees where each non-leaf node has at most one non-leaf child.
Such a tree is often called a “caterpillar” - because it’ll look like a line (the “body” of the caterpillar) with some leaves hanging off of each vertex on the line (the “legs” of the caterpillar).
Let the non-leaf vertices be x_1, x_2, x_3, \ldots, x_k, with x_1 = 1 and x_i being the parent of x_{i+1}.
As seen earlier (the case of the root having exactly one non-leaf child), when we assign values to these vertices, A_{x_1} \gt A_{x_2} \gt\ldots\gt A_{x_k} must hold.
Observe that if we’ve chosen these values:
- 0, 1, 2, \ldots, A_{x_k}-1 must all lie in the subtree of x_k; which in turn means they must appear among the children of x_k.
- 0, 1, 2, \ldots, A_{x_{k-1}}-1 must all lie in the subtree of x_{k-1}.
However, we already know 0, 1, 2, \ldots, A_{x_k} appear in its subtree (since they’re in the subtree of x_k, which is a child of x_{k-1}).
This means only the values A_{x_k}+1, \ldots, A_{x_{k-1}}-1 need to be assigned - and each of these must appear either at a leaf child of x_k or of x_{k-1} (since those are the only free vertices in the subtree of x_{k-1}). - In general, 0, 1, \ldots, A_{x_{i+1}} will appear in the subtree of x_{i+1}, so we’ll then need to assign A_{x_{i+1}}+1, \ldots, A_{x_i}-1 to some empty leaves in the subtree of x_i (which need not be leaves that are its direct children - descendant leaves are fine too).
In particular, note that the calculation for A_{x_i} depends only on A_{x_{i+1}}.
This allows us to formulate a dynamic programming solution.
Define dp(i, y) to be the number of ways of assigning the values 0, 1, 2, \ldots, y to the subtree of x_i, such that A_{x_i} = y.
That is, dp(i, y) denotes the number of configurations that have A_{x_i} = y, while ignoring all values larger than y.
To compute this: A_{x_i} = y, so A_{x_{i+1}} must be something smaller than y.
Suppose we fix A_{x_{i+1}} = z.
Then,
- All of 0, 1, 2, \ldots, z must appear in the subtree of x_{i+1}.
There are dp(i+1, z) ways of assigning these values. - That leaves the values z+1, \ldots, y-1 yet to be assigned.
These can each be assigned to any leaf in the subtree of x_i. - So, let s_{x_i} denote the subtree size of x_i (excluding x_i itself).
Of these s_{x_i} nodes, exactly z+1 of them have their values fixed already - and the unfilled values are surely all leaves. - That means there are (s_{x_i} - z - 1) empty leaves, and we want to assign the values z+1, \ldots, y-1 to some of them.
There are \binom{s_{x_i}-z-1}{y-z-1} ways to choose which leaves will contain these values, and then (y-z-1)! ways to arrange them among these leaves. - So, we obtain
This allows for dp(i, y) to be computed in linear time.
There are \mathcal{O}(N^2) states, so all the dp values can be computed in \mathcal{O}(N^3) time.
To compute the overall answer: note that dp(1, y) counts the number of ways to assign the values 0, 1, 2, \ldots, y to the vertices while having A_1 = y, and not caring about values \gt y.
However, once we’ve reached the root, values \gt y have no further constraints.
This means all remaining values can be freely distributed among the remaining leaves; and there are (N-1-y)! ways of doing so.
This makes the final answer be
TIME COMPLEXITY:
\mathcal{O}(N^3) per testcase.
CODE:
Editorialist's code (PyPy3)
mod = 998244353
fac = [1] +list(range(1, 1005))
for i in range(1, 1005): fac[i] = fac[i-1] * i % mod
C = [ [0 for i in range(1005)] for i in range(1005) ]
for n in range(1005):
C[n][0] = 1
for r in range(1, n+1):
C[n][r] = (C[n-1][r] + C[n-1][r-1]) % mod
for _ in range(int(input())):
n = int(input())
p = [0, 0] + list(map(int, input().split()))
leaf = [1]*(n+1)
for x in p: leaf[x] = 0
nxt = [-1]*(n+1)
good = 1
for i in range(2, n+1):
if leaf[i] == 0:
if nxt[p[i]] != -1: good = 0
nxt[p[i]] = i
if good == False:
print(0)
continue
subsz = [1]*(n+1)
for i in reversed(range(1, n+1)):
subsz[p[i]] += subsz[i]
dp = [ [0 for i in range(n+1)] for i in range(n+1) ]
for i in reversed(range(1, n+1)):
if leaf[i] == 1: continue
for j in range(subsz[i]):
if nxt[i] == -1:
dp[i][j] = C[subsz[i]-1][j] * fac[j]
else:
for k in range(j):
dp[i][j] = (dp[i][j] + dp[nxt[i]][k] * C[subsz[i]-k-2][j-k-1] * fac[j-k-1]) % mod
ans = 0
for i in range(n):
ans += dp[1][i] * fac[n-1-i] % mod
print(ans % mod)