2019湖南省赛 2019（树形DP,换根）

2019
Bobo 有一颗 $n$ 个点的树，点的编号是 $1, 2, \dots, n$. 树有 $(n - 1)$ 条边，第 $i$ 条边的端点是 $a_i$ 和 $b_i$，权值是 $c_i$. 求满足 $u < v$ 的 $(u, v)$ 数量，满足点 $u$ 到点 $v$ 路径上的权值和是 $2019$ 的倍数。

输入格式
输入文件包含多组数据，请处理到文件结束。

每组数据的第一行包含一个整数 $n$.

接下来 $(n - 1)$ 行，其中第 $i$ 行包含三个整数 $a_i$, $b_i$ 和 $c_i$.

$n \leq 2 \times 10^4$
$1 \leq a_i, b_i \leq n$
$0 \leq c_i < 2019$
$n$ 的总和不超过 $10^5$.
输出格式
对于每组数据，输出一个整数，表示所求的值。

思路：
又是关于树上点间距离的题目，因为每个点的情况都需要考虑，所以很明显是需要树形dp换根了，说人话就是要两个dfs。

定义以 $dp[i][j]$ 以 i 为根节点时，其他点到这个点距离mod2019值为j的数目。特别的 $dp[i][0] = 1;$

第一次递归是有根树情况下计算的结果。
这时候直接 $dp[u][(j + w) \mod 2019] += dp[v][j]$

第二次递归是无根树情况下计算的结果。
这次计算的时候还要加上从父节点来的数目。
$f[v][j] = dp[u][j - w] - dp[v][j-w*2] + dp[v][j]$

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <list>

using namespace std;

typedef long long ll;
const int maxn = 2e4 + 7;

int head[maxn],nex[maxn * 2],to[maxn * 2],val[maxn * 2],tot;
int dp[maxn][2019],a[maxn];
int n;
ll ans;

void init()
{
    for(int i = 1;i <= n;i++)
    {
        dp[i][0] = 1;
        for(int j = 1;j < 2019;j++)
        {
            dp[i][j] = 0;
        }
    }
    memset(head,0,sizeof(head));
    tot = 0;
    ans = 0;
}

void add(int x,int y,int z)
{
    to[++tot] = y;
    val[tot] = z;
    nex[tot] = head[x];
    head[x] = tot;
}

void DP1(int u,int fa)
{
    for(int i = head[u];i;i = nex[i])
    {
        int v = to[i],w = val[i];
        if(v == fa)continue;
        DP1(v,u);
        for(int j = 0;j < 2019;j++)
        {
            int k = (j + w) % 2019;
            dp[u][k] += dp[v][j];
        }
    }
}

void DP2(int u,int fa)
{
    for(int i = head[u];i;i = nex[i])
    {
        int v = to[i],w = val[i];
        if(v == fa)continue;
        for(int j = 0;j < 2019;j++)
        {
            a[j] = dp[v][j];
        }
        for(int j = 0;j < 2019;j++)
        {
            int k1 = (j + w) % 2019;
            int k2 = ((j - w) % 2019 + 2019) % 2019;
            dp[v][k1] += dp[u][j] - a[k2];
        }
        DP2(v,u);
    }
}

int main()
{
    while(~scanf("%d",&n))
    {
        init();
        for(int i = 1;i < n;i++)
        {
            int x,y,z;scanf("%d%d%d",&x,&y,&z);
            add(x,y,z);add(y,x,z);
        }
        DP1(1,-1);
        DP2(1,-1);
        for(int i = 1;i <= n;i++)
        {
            ans += dp[i][0] - 1;
        }
        printf("%lld\n",ans / 2);
    }
    return 0;
}