EagleBear2002 的博客

这里必须根绝一切犹豫,这里任何怯懦都无济于事

P5860 「SWTR-3」Counting Trees

题目背景

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

看着满树林的树,小 $\mathrm{S}$ 灵感一闪,想到了一道题目。

$$ \mathrm{Suddenly,\ Little\ S\ thought\ of\ a\ supercalifragilisticexpialidocious\ problem.} \ \mathrm{He\ wanted\ Little\ A\ to\ answer\ it.} $$

题目描述

现在有 $n$ 个点,每个点有一个权值 $v_i$。

小 $\mathrm{S}$ 想要小 $\mathrm{A}$ 从中选一些点组成一个集合,设集合 $S=\set{d_1,d_2,\dots,d_m}(1\leq m\leq n)$。

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

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

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

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

输入格式

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

第二行,$n$ 个整数 $v_1,v_2,\dots,v_n$

输出格式

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

输入输出样例 #1

输入 #1

1
2
3
1 1 1

输出 #1

1
3

输入输出样例 #2

输入 #2

1
2
5
1 2 1 3 1

输出 #2

1
8

输入输出样例 #3

输入 #3

1
2
8
1 2 1 2 4 1 3 1

输出 #3

1
44

输入输出样例 #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^{2}_{3}=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$。

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

对于前 $4$ 个 $\mathrm{Subtask}$,时间限制 $1\mathrm{s}$。

对于第 $5$ 个 $\mathrm{Subtask}$,时间限制 $3\mathrm{s}$。

对于后 $3$ 个 $\mathrm{Subtask}$,时间限制 $6\mathrm{s}$。

对于所有测试点,空间限制 $256\mathrm{MB}$。

题解

一个集合 $\set{d_1, d_2, \dots, d_m}$ 能组成树的充要条件是 $\sum_{i=1}^{m} v_{d_i} = 2m - 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{align*} G(x) & = \prod_{i=1}^{n}\left(1+x^{v_{i}-2}\right)[v_i > 2] \ & = \exp\left(\sum_{i=1}^{n}\ln(1+x^{a_{i}}) [v_i > 2] \right) \end{align*} $$

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

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