「Daily OI Round 1」block 题解

teylnol_evteyl · 2023-08-21 11:25:01 · 题解

设 c_u 表示节点 u 的颜色，f_u 表示只考虑原树中 u 的子树中的点、选择点 u 的方案数。对于儿子 v，可以选择同色儿子节点，也可以跳过这个儿子节点，选择 v 的与 u 同色的儿子节点 w，故状态转移方程为：

其中 v 是 u 的儿子，w 是 v 的儿子。

统计答案时，除了考虑选择的点的最近公共祖先被选择的情况（即 \sum f_u），还有最近公共祖先没有被选择的情况，也就是说一个节点 u 的几个儿子节点颜色相同，则可以分别在它们的子树中选择，而不选 u。设 g_{u,c} 表示考虑 u 的子树，不选择 u，选择至少两个颜色为 c 的儿子节点的方案数，因为要分别减去选择 0 或 1 个儿子的情况，所以：

其中 v 是 u 的儿子。

#include <bits/stdc++.h>

using namespace std;

typedef long long LL;

const int N = 1e5 + 5, P = 1e9 + 7;

int n, c[N];
int la[N], ne[N * 2], en[N * 2], idx;
LL f[N], g[N], res;

void add(int a, int b)
{
    ne[ ++ idx] = la[a];
    la[a] = idx;
    en[idx] = b;
}
void dp(int u, int fa)
{
    f[u] = 1;
    for(int i = la[u]; i; i = ne[i])
    {
        int v = en[i];
        if(v == fa) continue ;
        dp(v, u);

        LL t = 1;
        for(int j = la[v]; j; j = ne[j])
        {
            int w = en[j];
            if(w == u) continue ;
            if(c[w] == c[u]) t = t * (f[w] + 1) % P;
        }
        if(c[v] == c[u]) t = (t + f[v]) % P;
        f[u] = f[u] * t % P;
    }
    for(int i = la[u]; i; i = ne[i])
    {
        int v = en[i];
        if(v == fa) continue ;
        g[c[v]] = g[c[v]] * (f[v] + 1) % P;
    }
    for(int i = la[u]; i; i = ne[i])
    {
        int v = en[i];
        if(v == fa) continue ;
        res = (res + g[c[v]] - 1) % P;
        g[c[v]] = 1;
    }
}

int main()
{
    scanf("%d", &n);
    for(int i = 1; i <= n; i ++ ) scanf("%d", &c[i]);
    for(int i = 1; i < n; i ++ )
    {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    add(0, 1);
    for(int i = 1; i <= n; i ++ ) g[i] = 1;
    dp(0, 0);
    printf("%lld\n", res);
    return 0;
}