:::info[Description] Link. :::
:::infoSolution 先考虑一个暴力的 dp 做法。设 为：
- 若 则 表示只考虑 为根的子树，包含 且点集中只有颜色 的方案数。
- 若 则 表示只考虑 为根的子树，包含 且点集中恰好出现 两个颜色的方案数。
因此以 为最高点的答案就是 。
然后考虑转移。初始条件显然是 为根的子树中只选择 自己，因此就是 ，其余位置的 dp 值都为 。
转移考虑合并同色子结点，考虑当前处理的子树为 ，要合并 的某个子树 ，且有 。则先考虑 怎么算。这里需要分类讨论：
- 若父亲中原来只有颜色 的方案，则要么不选 子树，要么接上 中一个只有颜色 的连通块。因此有 。
- 若父亲中还有另一种颜色 ，则需要进一步分类讨论。此时不管是 这一侧还是 这一侧都要么是全 要么是既有 又有 。然后还有一种特殊情况是 子树内一个也不选。省略一些 dirty work 之后可得最终的转移方程形如 。
然后再考虑合并异色儿子的情况。设 。则如果此时选择了 子树内的点，因为 结点本身的颜色为 而最终总共只能有两种颜色，因此这个点集内的颜色只能是 。因此在 子树中只能选择只有颜色 即 或同时包含 两个颜色 两种情况。然后父亲结点这一侧也同样要么只有 即 ，要么同时有 两种颜色即 。因此新增的方案数为 。转移就是 ，其他 信息的值都不会发生任何变化。
最后的答案就是 。
然后很显然这个 dp 转移过程同样可以用线段树合并来优化，具体细节这里就不赘述了。 :::
:::success[Code]
namespace lowspeed_song {
inline void init() {
}
int a[N];
vector<int> adj[N];
struct Node {
int l, r, sum, tag;
inline void push(int v) { sum = sum * v % mod, tag = tag * v % mod; }
} tree[N << 6];
int idx = 0;
inline int newnode() { return tree[++idx].tag = 1, idx; }
inline void pushup(int rt) { tree[rt].sum = (tree[tree[rt].l].sum + tree[tree[rt].r].sum) % mod; }
inline void pushdown(int rt) {
if (tree[rt].tag != 1) tree[tree[rt].l].push(tree[rt].tag), tree[tree[rt].r].push(tree[rt].tag), tree[rt].tag = 1;
}
inline void mdf(int l, int r, int rt, int pos, int v) {
if (l == r) {
tree[rt].sum = (tree[rt].sum + v) % mod;
return;
} int mid = l + r >> 1; pushdown(rt);
if (pos <= mid) {
if (!tree[rt].l) tree[rt].l = newnode();
mdf(l, mid, tree[rt].l, pos, v);
} else {
if (!tree[rt].r) tree[rt].r = newnode();
mdf(mid + 1, r, tree[rt].r, pos, v);
} pushup(rt);
}
inline int qry(int l, int r, int rt, int pos) {
if (!rt) return 0;
if (l == r) return tree[rt].sum;
int mid = l + r >> 1; pushdown(rt);
if (pos <= mid) return qry(l, mid, tree[rt].l, pos);
return qry(mid + 1, r, tree[rt].r, pos);
}
inline void merge(int rt1, int rt2, int l, int r, int fu, int fv, int col) {
if (!rt1) return;
if (!rt2) return tree[rt1].push(fv + 1);
if (l == r) {
if (l == col) tree[rt1].sum = tree[rt1].sum * (tree[rt2].sum + 1) % mod;
else tree[rt1].sum = (tree[rt1].sum + tree[rt1].sum * fv % mod + tree[rt1].sum * tree[rt2].sum % mod + tree[rt2].sum * fu % mod) % mod;
return;
} pushdown(rt1), pushdown(rt2);
int mid = l + r >> 1;
if (tree[rt1].l && tree[rt2].l) merge(tree[rt1].l, tree[rt2].l, l, mid, fu, fv, col);
else if (tree[rt1].l) tree[tree[rt1].l].push(fv + 1);
else if (tree[rt2].l) tree[rt1].l = tree[rt2].l, tree[tree[rt1].l].push(fu);
if (tree[rt1].r && tree[rt2].r) merge(tree[rt1].r, tree[rt2].r, mid + 1, r, fu, fv, col);
else if (tree[rt1].r) tree[tree[rt1].r].push(fv + 1);
else if (tree[rt2].r) tree[rt1].r = tree[rt2].r, tree[tree[rt1].r].push(fu);
pushup(rt1);
}
int n, res;
inline void dfs(int u, int fa) {
tree[u].tag = 1, mdf(1, n, u, a[u], 1);
for (int &v : adj[u]) if (v != fa) {
dfs(v, u);
if (a[u] == a[v]) merge(u, v, 1, n, qry(1, n, u, a[u]), qry(1, n, v, a[v]), a[u]);
else mdf(1, n, u, a[v], (qry(1, n, u, a[u]) + qry(1, n, u, a[v])) * (qry(1, n, v, a[u]) + qry(1, n, v, a[v])) % mod);
} res = (res + tree[u].sum) % mod;
}
inline void sol([[maybe_unused]]int __testcase_id) {
cin >> n;
for (int i = 1; i <= n; ++i) cin >> a[i];
for (int i = 1; i < n; ++i) {
int a, b; cin >> a >> b;
adj[a].emplace_back(b), adj[b].emplace_back(a);
} idx = n, dfs(1, 0), cout << res << '\n';
}
} // namespace lowspeed_song:::
难度：黑-
鲜活指数：7
喜欢的话，留下你的评论吧～