牛客周赛8 D题

D-小美的树上染色_牛客周赛 Round 8 (nowcoder.com)

思路：

树形dp。每个点可以分为涂红（记为1）和不涂红（记为0）两种状态。将点1看做根节点进行dfs。

首先考虑根节点为0的情况：子节点是否为红色对它没有任何影响，且由于这是一棵树，一个点可以有许多子树，所以 f[u][0] 的转移方程是最大值的和。f[u][1]则是在发现可以和自己一起变红的子节点后，减去f[u][0]中加上的这个子节点为根的子树的点数（max( f[j][0], f[j][1] ) ），加上f[j][0](因为此时的子节点只能是白色），再加上2。

代码：

#include 
using namespace std;
typedef long long ll;
const int N=2e5+10;
int e[N],ne[N],h[N],w[N];
int idx;
int f[N][2];

void add(int a,int b){
    e[idx]=b;ne[idx]=h[a];h[a]=idx++;
}
bool dic(int x,int y){
    ll sum=x*y*1ll;
    ll half=sqrt(sum);
    if(half*half==sum)return true;
    return false;
}
void dfs(int u,int p){
    for(int i=h[u];~i;i=ne[i]){
        int j=e[i];
        if(j==p)continue;
        dfs(j,u);
        f[u][0]+=max(f[j][0],f[j][1]);
    }
    for(int i=h[u];~i;i=ne[i]){
        int j=e[i];
        if(j==p)continue;
        if(dic(w[u],w[j]))
            f[u][1]=max(f[u][1],f[u][0]-max(f[j][0],f[j][1])+f[j][0]+2);
    }
}
int main(){
    int n;
    cin>>n;
    memset(h,-1,sizeof h);
    for(int i=1;i<=n;i++){
        cin>>w[i];
    }
    for(int i=1;i>a>>b;
        add(a,b);
        add(b,a);
    }
    dfs(1,0);
    cout<

思考：

写dp的时候忘记掉这是一颗树了，把它当成一条链来写了。链式前向星储存树别忘了数组要开2倍，毕竟每条边得存两次。