hdu 4303 树形DP

转自：点击打开链接

献上我的膝盖

题意：给你一颗树，每个节点有权值，每条边有颜色，问所有满足相邻边颜色不同的路径的权值之和

分析：一开始看成求权值最大的路径，所以就是水题了，不过加个统计就得略想一下子，f[ i ]表示i的所有子节点的连到i 路径和，这里有个限制就是i到子节点的颜色与i到父节点的颜色不同，g[ i ]记录相应的边的数量，于是有f[ i ]=sum{ f[ j ]+g[ j ]*v[i] },  g[ i ]=sum{ g[ j ] }（ i和j满足颜色限制）这样，对于每个节点i有3种路径，一种以它为端点的路径，这个答案ans1=sum{ f[ j ]+g[ j ]*v[i] }这里没有颜色限制，第二种就是通过i到i的父节点，这个可以转化成父节点的问题，所以不用考虑，回溯的时候自然解决了，第三种最复杂，就是两个来自子节点的边在i点和成一条路径，这里因为需要颜色限制所以比较烦，不过，还是有简单的方法的，sumf = sum{ f[ j ]+g[ j ]*v[ i ] }, sumg=sum{ g[j ] }这里i和j没有颜色限制，对于来自每个子节点j的边我们这样统计，ans2=(sumf - sum{ f[ k ]+g[ k ]*v[i] })*g[ j ] + (sumg-sum{ g[k ] }) *f[ j ]   这里k表示子节点k到i的边与j到i的边颜色相同，至于为什么这样算自己想想就明白了，很简单，所以答案就是每个节点的ans1+ans2的和

注意使用scanf读入，cin超时啊T_T

AC代码如下：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <map>
using namespace std;

struct Edge{
    int to, next, colour;
};

long long f[310000], g[310000];
Edge edge[1000000];
int head[310000], tot;
int N, v[310000];
long long ans;
bool mark[310000];

void add_edge( int a, int b, int cc ){
    edge[tot].to = b;
    edge[tot].colour = cc;
    edge[tot].next = head[a];
    head[a] = tot++;
    edge[tot].to = a;
    edge[tot].colour = cc;
    edge[tot].next = head[b];
    head[b] = tot++;
}

void DFS( int pos, int cc ){

    mark[pos] = true;
    f[pos] = v[pos];
    g[pos] = 1;

    long long tempans1, tempans2, sumf, sumg;
    map<int,long long> mg, mf;

    tempans1 = tempans2 = 0;
    sumf = sumg = 0;

    for( int i = head[pos]; i != -1; i = edge[i].next ){
        int to = edge[i].to;
        int tempc = edge[i].colour;
        if( !mark[to] ){
            DFS( to, tempc );
            tempans1 += f[to] + g[to] * v[pos];
            sumf += f[to] + g[to] * v[pos];
            sumg += g[to];

            if( mg.find( tempc ) == mg.end() ){
                mg[tempc] = g[to];
                mf[tempc] = f[to] + g[to] * v[pos];
            }else{
                mg[tempc] += g[to];
                mf[tempc] += f[to] + g[to] * v[pos];
            }

            if( edge[i].colour != cc ){
                f[pos] += f[to] + g[to] * v[pos];
                g[pos] += g[to];
            }
        }
    }

    for( int i = head[pos]; i != -1; i = edge[i].next ){
        int to = edge[i].to;
        int tempc = edge[i].colour;
        if( !mark[to] ){
            tempans2 += ( sumf - mf[tempc] ) * g[to] + ( sumg - mg[tempc] ) * f[to];
        }
    }
    ans += tempans1 + tempans2 / 2;
    mark[pos] = false;
}

int main(){

    while( scanf( "%d", &N ) != EOF ){
        memset( head, -1, sizeof( head ) );
        tot = 0;
        for( int i = 1; i <= N; i++ ){
            scanf( "%d", &v[i] );
        }
        for( int i = 1; i < N; i++ ){
            int temp1, temp2, temp3;
            scanf( "%d%d%d", &temp1, &temp2, &temp3 );
            add_edge( temp1, temp2, temp3 );
        }
        memset( mark, false, sizeof( mark ) );
        ans = 0;
        DFS( 1, -1 );
        cout << ans << endl;
    }

    return 0;
}