bzoj 2152 聪聪可可 树形dp

题目链接

题意

给定一棵(n)个点的树， 每条边上有权值。在(n^2)个点对中，问有多少对点((u,v))，(u)与(v)之间所有边上数的和加起来恰好是(3)的倍数。

思路

对每个点，用一个三元组记录它的子树中距它的距离(\%3)为(0,1,2)的点的个数。

树形(dp)，从下往上计算+转移。

具体计算上需要注意：

1. 假设点(u)的三元组为((a,b,c))，经过点(u)的答案并不是(a*a+b*c+c*b)，这是因为会和之前的计算有重复（转移自(u)的同一个孩子的两个点，它们间的路径不过经过点(u)）。故需在转移之前进行计算。
2. 计算的原始思路是二重循环枚举(u)的两个孩子，累乘再累加，然而这样显然会(T)，怎么办呢？采取边计算边转移的方法——每次将当前(u)上记录的数据（即(u)本身+处理过了的(u)的所有孩子） 与 当前处理的孩子(v)的数据 做相乘处理，然后将(v)的数据也转移到(u)上，道理还是挺显然的。

Code

#include <bits/stdc++.h>
#define maxn 20010
using namespace std;
typedef long long LL;
LL gcd(LL a, LL b) { return b ? gcd(b, a % b) : a; }
int tot, ne[maxn];
LL ans, dis[maxn][3];
struct Edge {
    int to, ne, w;
    Edge(int _to=0, int _ne=0, int _w=0) : to(_to), ne(_ne), w(_w) {}
}edge[maxn * 2];
void add(int u, int v, int w) {
    edge[tot] = Edge(v, ne[u], w);
    ne[u] = tot++;
}
void dfs(int u, int fa) {
    dis[u][0] = 1;
    for (int i = ne[u]; ~i; i = edge[i].ne) {
        int v = edge[i].to;
        if (v == fa) continue;
        dfs(v, u);
        int w = edge[i].w;
        for (int i = 0; i < 3; ++i) {
            ans += dis[u][i] * dis[v][(6-w-i)%3];
            dis[u][i] += dis[v][(i+3-w)%3];
        }
    }
}
int main() {
    int n;
    scanf("%d", &n);
    tot = 0; memset(ne, -1, sizeof(ne));
    for (int i = 1; i < n; ++i) {
        int u, v; int w;
        scanf("%d%d%d", &u, &v, &w);
        w %= 3;
        add(u, v, w); add(v, u, w);
    }
    dfs(1, -1);
    ans <<= 1;
    ans += n;
    LL d = gcd(ans, 1LL*n*n);
    printf("%lld/%lld
", ans/d, 1LL*n*n/d);
    return 0;
}