2019年icpc沈阳网络赛D题

ACM-ICPC 2019 沈阳赛区网络预赛 D题
题目链接

换根就是先DFS一遍之后,根据第一遍DFS得到的结果去处理第二遍DFS,第二遍DFS就是将每个点都看成一个根.
本题令$cnt1[i][j]$为以$root$为根,$i$的子树中距$i$的距离%3为$j$的点的数量.$dp[i][j]$则是i的子树中距$i$的距离模3为$j$的距离路径总和.
这两种状态可以用一遍DFS求得, $root$就是自己设的根.
第二遍DFS就是处理每个点为根的情况.
考虑$ cnt2[i][j] $ , 在以$root$为根的情况下, 非$i$的子节点$i$的距离%3为j的节点个数.
$dp1[i][j$ ,在以$root$为根的情况下，非$i$的子节点到$i$的距离%3为$j$的距离和。

令$v$为当前节点, $u$为其父节点

推出式子

$cnt2[v][t]=((cnt2[u][j]+cnt1[u][j]-cnt1[v][t1])\%mod+mod)\%mod;$
$dp1[v][t]=((dp[u][j]-w\ast cnt1[v][t1]+mod)\%mod+w\ast cnt2[v][t]-dp[v][t1]+mod+dp1[u][j])\%mod$
其中因为u的子节点包括了v的子节点, 需要处理掉.就是减掉.在加上父亲非子节点的贡献.
考虑路径和.同理, 减掉多余的贡献.加上父亲节点的非子节点的贡献.

/************************************************************]
> File Name: D.cpp
> Author: ccdxc
> Mail: linyi0430@gmail.com
> Created Time: 2019年09月14日 星期六 21时31分19秒
***************************************************************/

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#include<vector>
#include<queue>
#include<set>
#include<map>
#include<string>
#include<cmath>
#include<stdlib.h>
#include<time.h>
#define ll long long
#define pb push_back
using namespace std;
const int INF = 0x3f3f3f3f;
const int Maxn = 1e5 + 10;
const ll mod = 1e9 + 7;
int n , head[Maxn];
int cnt;
ll cnt1[Maxn][3] , dp[Maxn][3];
ll ans[3] , cnt2[Maxn][3] , dp2[Maxn][3];

struct node{
    int w ;
    int next , to;
}edge[Maxn];

//初始化
void init(){
    cnt = 1;
    memset(head ,  -1 , sizeof(head));
    ans[0] = ans[1] = ans[2] = 0;
    memset(dp,0,sizeof(dp));
    memset(dp2,0,sizeof(dp2));
    memset(cnt1,0,sizeof(cnt1));
    memset(cnt2 , 0 ,sizeof(cnt2));
}
//链式前向星存树
void add(int u, int v , ll w){
    edge[cnt].w = w;
    edge[cnt].to = v;
    edge[cnt].next = head[u];
    head[u] = cnt++;

}

void dfs1(int u,int fa){
    cnt1[u][0] = 1;
    for(int i = head[u] ;i != -1;i = edge[i].next){
        int v = edge[i].to , w = edge[i].w;
        if(v == fa) continue;
        dfs1(v , u);
        for(int j = 0;j < 3;j++){
            int t = (w + j) % 3;
            cnt1[u][t] = (cnt1[u][t] + cnt1[v][j])%mod;
            dp[u][t] = (dp[u][t] + dp[v][j] + w * cnt1[v][j]) %mod;
        }
    }
}

void dfs2(int u , int fa){
    for(int i = head[u];i != -1;i = edge[i].next){
        int v = edge[i].to , w = edge[i].w;
        if(v == fa) continue;
        for(int j = 0;j < 3;j++){
            int t = (w + j)%3;
            int t1 = ((j - w) % 3 + 3 ) % 3;
            cnt2[v][t] = ((cnt1[u][j] - cnt1[v][t1] + cnt2[u][j])%mod + mod )%mod;
            dp2[v][t] = ((dp[u][j] - dp[v][t1]-w*cnt1[v][t1]) % mod + dp2[u][j] + w * cnt2[v][t]) % mod;
        }
        dfs2(v , u);
    }
    
}

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    int u , v , w;
    while(cin >> n){
        init();
        for(int i = 1;i < n;i++){
            cin >> u >> v >> w;
            u++ , v++;
            add(u, v , w);
            add(v , u , w);
        }
        //第一遍dfs求已每个节点为子树时，每个子节点的距离 以及 模3的个数
        dfs1(1 , -1);
        //第二遍dfs，通过换根求出，每个节点与父亲节点的距离以及模3的个数
        dfs2(1,-1);
        //最后答案就是每个节点的父亲节点的距离加上每个节点的子节点之和，遍历一遍即可
        for(int i = 1;i <= n;i++){
            for(int j = 0;j < 3;j++){
                ans[j] =( ans[j]+ dp[i][j] + dp2[i][j] ) %mod;
            }
        }
        cout << ans[0] <<" " << ans[1] <<" "<< ans[2]<<endl;
    }
}
/*
6
0 1 5
1 2 2
2 3 1
1 4 4
4 5 3
*/