CF1146F - Leaf Partition
Solution
感觉做这种细节很多的分类讨论树形dp还是有点乱。
大概有一个 n a i v e naive naive的想法是,我们令 f x , 0 / 1 f_{x,0/1} fx,0/1表示以 x x x为根的子树, x x x结点有没有颜色的方案数。
然后如果存在两个有颜色的儿子,则 x x x的颜色可以直接确定。
但是直接这样不好转移(因为当前无色的点可能在更浅的地方和其他点连通导致有了颜色),我们修改状态为 f x , 0 / 1 / 2 f_{x,0/1/2} fx,0/1/2, 0 0 0表示 x x x无色, 1 1 1表示我们钦定了 x x x为一种它儿子的颜色, 2 2 2表示 x x x有一种确定的颜色。
那么
f
x
,
0
=
∏
v
f
v
,
0
+
f
v
,
2
f
x
,
1
=
(
∏
v
f
v
,
0
+
f
v
,
2
)
(
∑
v
f
v
,
2
f
v
,
0
+
f
v
,
2
)
+
(
∏
v
f
v
,
0
+
f
v
,
2
)
(
∑
v
f
v
,
1
f
v
,
0
+
f
v
,
2
)
f
x
,
2
=
(
∏
v
f
v
,
0
+
f
v
,
1
+
f
v
,
2
+
f
v
,
2
)
−
(
∏
v
f
v
,
0
+
f
v
,
2
)
−
(
∏
v
f
v
,
0
+
f
v
,
2
)
(
∑
v
f
v
,
1
+
f
v
,
2
f
v
,
0
+
f
v
,
2
)
f_{x,0}=\prod_v{f_{v,0}+f_{v,2}}\\ f_{x,1}=(\prod_v{f_{v,0}+f_{v,2})(\sum_v{\frac{f_{v,2}}{f_{v,0}+f_{v,2}}}})+(\prod_v{f_{v,0}+f_{v,2})(\sum_v{\frac{f_{v,1}}{f_{v,0}+f_{v,2}}}})\\ f_{x,2}=(\prod_v{f_{v,0}+f_{v,1}+f_{v,2}+f_{v,2})-(\prod_v{f_{v,0}+f_{v,2}})-(\prod_v{f_{v,0}+f_{v,2}})(\sum_v{\frac{f_{v,1}+f_{v,2}}{f_{v,0}+f_{v,2}}}})
fx,0=v∏fv,0+fv,2fx,1=(v∏fv,0+fv,2)(v∑fv,0+fv,2fv,2)+(v∏fv,0+fv,2)(v∑fv,0+fv,2fv,1)fx,2=(v∏fv,0+fv,1+fv,2+fv,2)−(v∏fv,0+fv,2)−(v∏fv,0+fv,2)(v∑fv,0+fv,2fv,1+fv,2)
时间复杂度 O ( n l g n ) O(nlgn) O(nlgn)(快速幂)。
Code
#include <bits/stdc++.h>
using namespace std;
template<typename T> inline bool upmin(T &x, T y) { return y <= x ? x = y, 1 : 0; }
template<typename T> inline bool upmax(T &x, T y) { return x <= y ? x = y, 1 : 0; }
#define MP(A,B) make_pair(A,B)
#define PB(A) push_back(A)
#define SIZE(A) ((int)A.size())
#define LEN(A) ((int)A.length())
#define FOR(i,a,b) for(int i=(a);i<(b);++i)
#define fi first
#define se second
typedef long long ll;
typedef unsigned long long ull;
typedef long double lod;
typedef pair<int, int> PR;
typedef vector<int> VI;
const lod eps = 1e-9;
const lod pi = acos(-1);
const int oo = 1 << 30;
const ll loo = (1ll << 62) - 1;
const int MAXN = 200005;
const int mods = 998244353;
const int MX = 100000;
const int inv2 = (mods + 1) >> 1;
const int INF = 0x3f3f3f3f; //1061109567
/*--------------------------------------------------------------------*/
namespace FastIO{
constexpr int SIZE = (1 << 21) + 1;
int num = 0, f;
char ibuf[SIZE], obuf[SIZE], que[65], *iS, *iT, *oS = obuf, *oT = obuf + SIZE - 1, c;
#define gc() (iS == iT ? (iT = ((iS = ibuf) + fread(ibuf, 1, SIZE, stdin)), (iS == iT ? EOF : *iS ++)) : *iS ++)
inline void flush() {
fwrite(obuf, 1, oS - obuf, stdout);
oS = obuf;
}
inline void putc(char c) {
*oS ++ = c;
if (oS == oT) flush();
}
inline void getc(char &c) {
for (c = gc(); !isalpha(c) && c != EOF; c = gc());
}
inline void reads(char *st) {
char c;
int n = 0;
getc(st[++ n]);
for (c = gc(); isalpha(c) ; c = gc()) st[++ n] = c;
st[++ n] = '\0';
}
template<class I>
inline void read(I &x) {
for (f = 1, c = gc(); c < '0' || c > '9' ; c = gc()) if (c == '-') f = -1;
for (x = 0; c >= '0' && c <= '9' ; c = gc()) x = (x << 3) + (x << 1) + (c & 15);
x *= f;
}
template<class I>
inline void print(I x) {
if (x < 0) putc('-'), x = -x;
if (!x) putc('0');
while (x) que[++ num] = x % 10 + 48, x /= 10;
while (num) putc(que[num --]);
}
inline void putstr(string st) {
for (int i = 0; i < (int)st.size() ; ++ i) putc(st[i]);
}
struct Flusher_{~Flusher_(){flush();}} io_Flusher_;
}
using FastIO :: read;
using FastIO :: putc;
using FastIO :: putstr;
using FastIO :: reads;
using FastIO :: print;
int f[MAXN][3];
vector<int> e[MAXN];
int upd(int x, int y) { return x + y >= mods ? x + y - mods : x + y; }
int quick_pow(int x, int y) {
int ret = 1;
for (; y ; y >>= 1) {
if (y & 1) ret = 1ll * ret * x % mods;
x = 1ll * x * x % mods;
}
return ret;
}
void tree_dp(int x) {
for (auto v : e[x]) tree_dp(v);
if (!e[x].size()) { f[x][0] = f[x][1] = 0, f[x][2] = 1; return; }
f[x][0] = 1;
int mul = 1, Mul = 1, sum = 0;
for (auto v : e[x]) {
f[x][0] = 1ll * f[x][0] * upd(f[v][0], f[v][2]) % mods;
f[x][1] = upd(f[x][1], 1ll * upd(f[v][1], f[v][2]) * quick_pow(upd(f[v][0], f[v][2]), mods - 2) % mods);
int t = upd(f[v][0], f[v][2]), invt = quick_pow(t, mods - 2);
mul = 1ll * mul * t % mods;
Mul = 1ll * Mul * upd(upd(f[v][0], f[v][1]), upd(f[v][2], f[v][2])) % mods;
sum = upd(sum, 1ll * invt * upd(f[v][1], f[v][2]) % mods);
}
f[x][1] = 1ll * f[x][1] * f[x][0] % mods;
f[x][2] = upd(Mul, mods - upd(mul, 1ll * mul * sum % mods));
// cout << x << ":" << e[x].size() << " " << mul << " " << Mul << " " << sum << " " << f[x][0] << " " << f[x][1] << " " << f[x][2] << endl;
}
signed main() {
#ifndef ONLINE_JUDGE
freopen("a.in", "r", stdin);
#endif
int n;
read(n);
for (int i = 2, x; i <= n ; ++ i) read(x), e[x].PB(i);
tree_dp(1);
print(upd(f[1][0], f[1][2]));
return 0;
}