2019 树形—DP

2019 树形—DP

题意：

给你一颗树，然你寻找 多少对（u, v) （u > v)的路径数 是2019的倍数。

题解：

树形dp入门专题

设 $dp[u][i] $ 以u为节点的子树中， 以u为起点 所有 路径之和 % 2019 等于i的个数。

那么答案的贡献有两种

1. 以(i为节点的子树 1 le i le n, dp[i][0])
2. 经过i节点 且属于 以i节点为根节点的子树。

第二种怎么算？

当我们在更新的 dp数组的时候这个时候就可算了， 类似与树的直径 dp的写法。

代码：

#include<bits/stdc++.h>
using namespace std;
const int N = 2e4 + 7;
typedef long long ll;
vector<pair<int, int> > g[N];

int n;
ll dp[N][2021], ans = 0;

void dfs(int u, int fa) {
    for (auto it: g[u]) {
        int to = it.first;
        int cost = it.second;
        if (to == fa) continue;
        dfs(to, u);
        ll cn = 0;
        for (int i = 0; i < 2019; i++) {
            if (dp[to][i]) {
                ll cat = dp[to][i];
                cn += dp[u][(2019 - (i + cost) % 2019 + 2019) % 2019] * cat;  
            }
        }
        cn += dp[u][(2019 - cost) % 2019];
        ans += cn;
        for (int i = 0; i < 2019; i++) {
            if (dp[to][i]) {
                int cnt = dp[to][i];
                dp[u][(i + cost) % 2019] += cnt; 
            }
        }
        dp[u][cost]++;
    }
}

int main() {
    while (~scanf("%d", &n)) {
        ans = 0;
        for (int i = 0; i <= n; i++) {
            g[i].clear();
        }
        for (int i = 0; i <= n; i++) {
            for (int j = 0; j < 2019; j++) {
                dp[i][j] = 0;
            }
        }
        for (int i = 1; i < n; i++) {
            int u, v, w;
            scanf("%d %d %d", &u, &v, &w);
            g[u].push_back({v, w});
            g[v].push_back({u, w});
        }
        dfs(1, 0);
        for (int i = 1; i <= n; i++) {
            ans += dp[i][0];
        }
        printf("%lld
", ans);
    }
}