P5860 「SWTR-3」Counting Trees

题目背景

一个风和日丽的早晨，小 $S$ 带着他的好朋友小 $A$ 在小树林里面数树。

看着满树林的树，小 $S$ 灵感一闪，想到了一道题目。

Suddenly, Little S thought of a supercalifragilisticexpialidocious problem .

题目描述

现在有 $n$ 个点，每个点有一个权值 $v_{i}$ 。

小 $S$ 想要小 $A$ 从中选一些点组成一个集合，设集合 $S = {d_{1}, d_{2}, \dots, d_{m}} (1 \leq m \leq n)$ 。

当然，小 $A$ 还需要保证这些点能形成一颗树，且 $d_{i}$ 的度数为 $v_{d_{i}} (i \in [1, m])$ 。（节点的度数：与它相邻的节点的个数）

小 $S$ 想问小 $A$ 有多少种满足条件的方案。

小 $A$ 深知自己肯定不会这道题目，所以他就拿来问你了。

由于方案数可能很大，所以请对 $998244353$ 取模。

输入格式

第一行，一个整数 $n$ 。

第二行， $n$ 个整数 $v_{1}, v_{2}, \dots, v_{n}$

输出格式

一行一个整数，表示方案数。

输入输出样例 #1

输入 #1

1
2

3
1 1 1

输出 #1

输入输出样例 #2

输入 #2

1 2	`5 1 2 1 3 1`

输出 #2

输入输出样例 #3

输入 #3

1 2	`8 1 2 1 2 4 1 3 1`

输出 #3

输入输出样例 #4

输入 #4

1 2	`50 8 1 10 2 2 1 2 1 1 2 5 1 11 6 13 13 10 4 1 13 11 2 2 11 13 10 1 1 4 3 4 2 15 2 2 1 1 2 1 7 14 2 2 4 13 2 7 5 6 10`

输出 #4

1	`176873472`

说明/提示

样例说明

对于样例 $1$ ，在三个节点中任选两个即可，答案为 $C_{3}^{2} = 3$ 。
对于样例 $2$ ，如图，共有 $8$ 种选择节点的方法。

数据范围与约定

本题使用捆绑测试。

Subtask 编号	$n \leq$	特殊性质	得分
$1$	$20$	无	$11$
$2$	$50$	无	$12$
$3$	$300$	无	$10$
$4$	$2500$	无	$17$
$5$	$4 \times 10^{4}$	无	$6$
$6$	$3 \times 10^{5}$	$v_{i} \leq 3$	$8$
$7$	$3 \times 10^{5}$	数据随机	$7$
$8$	$5 \times 10^{5}$	无	$29$

对于 $100 %$ 的数据， $2 \leq n \leq 5 \times 10^{5}$ ， $1 \leq v_{i} \leq n$ 。

$Subtask 7$ 中“数据随机”指：对于所有 $v_{i}$ ， $\frac{1}{3}$ 的概率为 1， $\frac{2}{3}$ 的概率为 $[2, n]$ 中等概率选择一个数。

对于前 $4$ 个 $Subtask$ ，时间限制 $1 s$ 。

对于第 $5$ 个 $Subtask$ ，时间限制 $3 s$ 。

对于后 $3$ 个 $Subtask$ ，时间限制 $6 s$ 。

对于所有测试点，空间限制 $256 MB$ 。

题解

一个集合 ${d_{1}, d_{2}, \dots, d_{m}}$ 能组成树的充要条件是 $\sum_{i = 1}^{m} v_{d_{i}} = 2 m - 2$ ，变形得到 $\sum_{i = 1}^{m} (v_{d_{i}} - 2) = - 2$ 。则问题转化为求 $v_{d_{i}} - 2$ 之和为 $- 2$ 的子集数。

构造生成函数

F (x) = \prod_{i = 1}^{n} (1 + x^{v_{i} - 2})

则答案为 $[x^{- 2}] F (x)$ ，即 $F (x)$ 中 $x^{- 2}$ 项的系数。

对于负整数次幂，考虑到 $v_{i} - 2 < 0$ 当且仅当 $v_{i} = 1$ ， $k$ 个 $v_{i} - 2 = - 1$ 其总贡献为 $(1 + x^{- 1})^{k} = \sum_{i = 0}^{k} C_{k}^{i} x^{- i}$ ；
对于 0 次幂，显然其贡献为 2 的幂；
对于正整数次幂，构造多项式

\begin{aligned} G (x) & =\prod_{i = 1}^{n}(1 + x^{v_{i} - 2})[v_{i}>2] \\ =\exp(\sum_{i = 1}^{n} \ln (1 + x^{a_{i}}) [v_{i} > 2]) \end{aligned}

其中 $\ln (1 + x^{p}) = \sum_{i = 1}^{+ \infty} \frac{(- 1)^{i - 1} x^{p i}}{i}$ 。

总时间复杂度为 $O (n \log n)$ 。