【树形DP+直径思想】代码源每日一题div1 三进制循环

三进制循环 - 题目 - Daimayuan Online Judge

题意：

 

思路：

有点像树的直径

回顾一下我们是怎么求直径的：维护根节点到其子树上的点的最大距离和次大距离，然后答案就是统计所有结点的次大值+最大值 的 最大值

 这道题也是一样的，只不过链有了顺序

一开始想维护一下最大值和次大值，后来想想根本不需要，两边都得最大值

注意到在统计答案的时候，需要统计从左到右的链的最大值，也要统计从右到左的链的最大值

因此需要记录结点到其子树的路径的方向，是往下的还是往上的

设dp[u][0]为，u到以u为根节点的子树中的往上走的路径的长度最大值

dp[u][1]为u到以u为根节点的子树中的往下走的路径的长度最大值

那么答案就是max{dp[u][0]+dp[u][1]}

转移非常好转移

如果是向上走的路径，判断一下(a[v]+1)%3+1=a[u]是否成立

如果是向下走的路径，判断一下a[v]=(a[u]+1)%3+1是否成立

Code：

#include <bits/stdc++.h>

#define low(x) (x&(-x))
#define int long long

using namespace std;

const int mxn=5e5+10;
const int mxe=5e5+10;

struct ty{
    int to,next;
}edge[mxe<<2];

int N,u,v,tot=0,ans=0;
int a[mxn],dp[mxn][2];//向上为0，向下为1
int head[mxn];

void add(int u,int v){
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void G_init(){
    tot=0;
    for(int i=0;i<=N;i++){
        head[i]=-1;
    }
}
void dfs(int u,int fa){
    dp[u][0]=dp[u][1]=1;
    for(int i=head[u];~i;i=edge[i].next){
        if(edge[i].to==fa) continue;
        dfs(edge[i].to,u);
        //处理向上：
        if((a[edge[i].to]+1)%3==a[u]) dp[u][0]=max(dp[u][0],dp[edge[i].to][0]+1);
        //处理向下
        if((a[u]+1)%3==a[edge[i].to]) dp[u][1]=max(dp[u][1],dp[edge[i].to][1]+1);
        ans=max(ans,dp[u][0]+dp[u][1]-1);
    }
}
void solve(){
    cin>>N;
    G_init();
    for(int i=1;i<=N-1;i++){
        cin>>u>>v;
        add(u,v);
        add(v,u);
    }
    for(int i=1;i<=N;i++) cin>>a[i];
    dfs(1,0);
    cout<<ans<<'\n';
}
signed main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    int __=1;//cin>>__;
    while(__--)solve();return 0;
}