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\n", ans/d, 1LL*n*n/d);
    return 0;
}