Black and Red vertices of Tree CodeChef - BLREDSET (树型dp套路题，可惜俺之前8会)

最新推荐文章于 2024-05-07 15:56:26 发布

sunmaoxiang

最新推荐文章于 2024-05-07 15:56:26 发布

阅读量279

点赞数 1

本文链接：https://blog.youkuaiyun.com/sunmaoxiang/article/details/102542526

版权

本文深入探讨了树形DP算法的应用，特别是在解决特定类型的计数问题中，如何通过树状结构进行有效递归分解，利用关键点判断和联通块统计，实现复杂度优化。文章详细解释了算法思路，包括关键点的定义、联通图的统计方法，以及如何通过累乘技巧高效计算以特定节点为根的联通图数量。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

题目大意

给定一棵N个结点的树，树上有三种颜色，0，1，2分别代表白，黑，红。
问题是有多少个联通块都为白色，且删掉这个联通块至少有一个红黑结点不相连。

思路

树上有一些关键点为白色且至少有两颗子树分别含有红黑两种结点，我们只需统计出这样的结点有哪些，然后求出其至少包含一个关键点的联通图有几种就行了，但是求”至少包含一个“不好求说实话，可以求出所有以r为根的白色联通图个数，减去以r为根不含有关键点的联通图个数即可。

求以u为根的联通图个数可以用累乘的方法，分别然后乘上u其它子树大小 + 1。

在这里插入图片描述

代码

#include <map>
#include <set>
#include <cmath>
#include <vector>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <iostream>
#include <algorithm>
using namespace std;

typedef long long ll;
const ll mod = 1e9 + 7;
const int maxn = 1e5 + 100;
int n;
vector<int> G[maxn];
int sz[maxn][3],col[maxn];
bool good[maxn];
ll tot[maxn], cntbad[maxn];
void dfs0(int u, int fa) {
    for(int i = 0; i < 3; i++) sz[u][i]=0;
    sz[u][col[u]] = 1;
    bool havered=false, haveblack=false;
    for(auto & v : G[u]) {
        if(v == fa) continue;
        dfs0(v,u);
        
        good[u] |= (havered && sz[v][1]);
        good[u] |= (haveblack && sz[v][2]);
        haveblack |= sz[v][1];
        havered |= sz[v][2];
        
        for(int i = 0; i < 3; i++) sz[u][i]+=sz[v][i];
    }
}
void dfs1(int u, int fa) {
    tot[u] = !col[u];
    cntbad[u] = (!col[u] && !good[u]);
    for(auto & v : G[u]) {
        if(v == fa) continue;
        dfs1(v, u);
        tot[u] = tot[u] * (tot[v] + 1) % mod;
        cntbad[u] = cntbad[u] * (cntbad[v] + 1) % mod;
    }
}
int main() {
    //freopen("/Users/maoxiangsun/MyRepertory/input.txt", "r", stdin);
    int T;
    cin >> T;
    while(T--) {
        cin>>n;
        for(int i = 1; i <= n; i++) {good[i]=false;G[i].clear();}
        for(int i = 0; i < n - 1; i++) {
            int u,v;cin>>u>>v;
            G[u].push_back(v);
            G[v].push_back(u);
        }
        for(int i = 1; i <= n; i++) cin>>col[i];
        
        dfs0(1, -1);
        for(int u = 1; u <= n; u++) {
            good[u] |= (sz[u][2] && sz[1][1] - sz[u][1]);
            good[u] |= (sz[u][1] && sz[1][2] - sz[u][2]);
        }
        dfs1(1, -1);
        ll ans = 0;
        for(int i = 1; i <= n; i++) {
            ans = (ans + tot[i] - cntbad[i] + mod) % mod;
        }
        printf("%lld\n",ans);
    }
    return 0;
}
/*
 2
 6
 1 2
 1 3
 1 4
 3 5
 3 6
 0 1 0 1 2 0
 6
 1 2
 1 3
 1 4
 3 5
 3 6
 1 0 0 0 2 0
 */