Gym 101353H 树型DP

题意:

求一棵树的所有简单路径权值之和,节点1为树的根

第一行T组样例

每个样例一个n(多少个节点)

接下来n-1行每行u,v,w(两个节点,两节点的距离(权值))

思路:

设dp[u]为以u为根的子树的价值,sz[u]为以u为根的子树的节点数,val[u]为u与所有子节点的距离之和
deep[u]为u的深度,设u为当前节点,v为u的子节点

val[u]=∑(val[v]+sz[v]*w(u,v))
设tot为val[u]中除去u到子树v的所有节点的距离,即tot=val[u]-(val[v]+sz[v]*w(u,v))
考虑每个子树v对dp[u]的贡献:dp[v]+sz[v]*tot(因为每条路径都可以在子树v中有sz[v]种不同的延伸)//路径是会跨过根的,根不一定是起点
最后得到:dp[u]=∑(dp[v]+sz[v]*tot)+val[u]

代码:

#include <iostream>
#include <stdio.h>
#include <algorithm>
#include <math.h>
#include <map>
#include <vector>
#include <string.h>
using namespace std;

typedef long long LL;
const int N = 1e5+5;
const LL mod = 1e9+7;
struct Node {
    int v,next;
    LL w;
}G[N<<1];
int head[N],cnt;

void init()
{
    memset(head, -1, sizeof(head));
    cnt = 0;
}

void add(int u, int v, LL w)
{
    G[cnt].v = v;
    G[cnt].w = w;
    G[cnt].next = head[u];
    head[u] = cnt++;
}
LL dp[N], sz[N], val[N];
int dep[N];
void DP(int u, int fa, int deep)
{
    dp[u]=0;sz[u]=1;val[u]=0;dep[u]=deep;
    for(int i=head[u];~i;i=G[i].next) {
        int v = G[i].v;
        LL w = G[i].w;
        if(v==fa) continue;
        DP(v,u,deep+1);
        val[u]=(val[u]+val[v]+sz[v]*w)%mod;
        sz[u]+=sz[v];
    }
    dp[u]=(dp[u]+val[u])%mod;
    for(int i = head[u];~i; i=G[i].next) {
        int v = G[i].v;
        LL w = G[i].w;
        if(v==fa) continue;
        LL tot=val[u]-val[v]+mod;
        LL cnt1=w*sz[v]%mod;
        dp[u]=(dp[u]+(tot-cnt1+mod)*sz[v])%mod;
    }
    return;
}

int main()
{
    int T, n;
    scanf("%d", &T);
    for(int cas = 1; cas <= T; cas++) {
        init();
        scanf("%d", &n);
        for(int i = 1; i < n; i++) {
            int u,v;
            LL w;
            scanf("%d%d%I64d", &u, &v, & w);
            add(u,v,w);
            add(v,u,w);
        }
        DP(1, -1, 1);
        LL ans = 0;
        for(int i = 1; i <= n; i++) ans=(ans+dp[i]*dep[i])%mod;
        printf("Case %d: %I64d\n", cas, ans);
    }
    return 0;
}





猜你喜欢

转载自blog.csdn.net/qihang_qihang/article/details/79876285
今日推荐