2019长安大学ACM校赛网络同步赛C LaTale （树上DP）

链接：https://ac.nowcoder.com/acm/contest/897/C
来源：牛客网

LaTale

时间限制：C/C++ 2秒，其他语言4秒
空间限制：C/C++ 32768K，其他语言65536K
64bit IO Format: %lld

题目描述

    Legend goes that in the heart of ocean, exists a gorgeous island called LaTale, which has n cities. Specially, there is only one way between every two cities.

    In other words, n cities construct a tree connected by n-1 edges. Each of the edges has a weight w.

Define d(u, v) the length between city u and city v. Under the condition of u differing from v, please answer how many pairs(u, v) in which d(u, v) can be divisible by 3.
    Pair(u,v) and pair(v,u) are considered the same.

输入描述:

    The first line contains an integer number T, the number of test cases.

    For each test case:

    The first line contains an integer n($\mathrm{2\le n\le 1052\le n\le 105),\; the\; number\; of\; cities.}$

    The following n-1 lines, each contains three integers ${\mathrm{uiui,\; vivi,\; wiwi(1\le ui,vi\le n,1\le wi\le 10001\le ui,vi\le n,1\le wi\le 1000),\; the\; edge\; of\; cities.}}^{}$

输出描述:

For each test case print the number of pairs required.

示例1

输入

复制

2
4
1 2 1
2 3 2
2 4 2
4
1 2 3
2 3 3
2 4 3

输出

复制

2
6

题意:
给你一个含有n个节点的树，
问有多少对节点u，v，他们的距离dist(u,v)%3==0

思路：
树上dp
我们定义状态 dp[i][0/1/2] 
表示第i个节点的子树中，距离第i个节点的距离%3的分别等于0、1、2三种情况的节点个数。
然后根据节点之间的关系转移。
对于一堆节点u，v。对答案的贡献是：
dp[u][j]*dp[v][(3-w-j+3)%3];
j 分别取0 1 2
然后再把底层的节点信息转移给它的父节点即可、

细节见代码：

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <vector>
#include <iomanip>
#define ALL(x) (x).begin(), (x).end()
#define rt return
#define dll(x) scanf("%I64d",&x)
#define xll(x) printf("%I64d\n",x)
#define sz(a) int(a.size())
#define all(a) a.begin(), a.end()
#define rep(i,x,n) for(int i=x;i<n;i++)
#define repd(i,x,n) for(int i=x;i<=n;i++)
#define pii pair<int,int>
#define pll pair<long long ,long long>
#define gbtb ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)
#define MS0(X) memset((X), 0, sizeof((X)))
#define MSC0(X) memset((X), '\0', sizeof((X)))
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define eps 1e-6
#define gg(x) getInt(&x)
#define db(x) cout<<"== [ "<<x<<" ] =="<<endl;
using namespace std;
typedef long long ll;
ll gcd(ll a, ll b) {return b ? gcd(b, a % b) : a;}
ll lcm(ll a, ll b) {return a / gcd(a, b) * b;}
ll powmod(ll a, ll b, ll MOD) {ll ans = 1; while (b) {if (b % 2)ans = ans * a % MOD; a = a * a % MOD; b /= 2;} return ans;}
inline void getInt(int* p);
const int maxn = 100010;
const int inf = 0x3f3f3f3f;
/*** TEMPLATE CODE * * STARTS HERE ***/
struct edge
{
    int next;
    int to;
    int dis;
    edge()
    {

    }
    edge(int nn, int dd)
    {
        next = nn;
        dis = dd;
    }
};
edge e[maxn << 1];
ll ans = 0ll;
int tot;
int head[maxn << 1];
void addedge(int a, int b, int c)
{
    e[tot].to = b;
    e[tot].dis = c;
    e[tot].next = head[a];
    head[a] = tot++;
}
void init()
{
    tot = 1;
}
ll dp[maxn][3];

void dfs(int id, int pre)
{
    dp[id][0] = 1ll;
    for (int i = head[id]; i != 0; i = e[i].next)
    {
        edge x = e[i];
        ll w = x.dis;
        if (x.to != pre)
        {
            dfs(x.to, id);
            ans += dp[id][0] * dp[x.to][(3 - w - 0 + 3) % 3];
            ans += dp[id][1] * dp[x.to][(3 - w - 1 + 3) % 3];
            ans += dp[id][2] * dp[x.to][(3 - w - 2 + 3) % 3];
            dp[id][(0 + w) % 3] += dp[x.to][0];
            dp[id][(1 + w) % 3] += dp[x.to][1];
            dp[id][(2 + w) % 3] += dp[x.to][2];

        }
    }

}
int main()
{
    //freopen("D:\\common_text\\code_stream\\in.txt","r",stdin);
    //freopen("D:\\common_text\\code_stream\\out.txt","w",stdout);

    int t;
    gbtb;
    cin >> t;
    while (t--)
    {
        ans = 0ll;
        int n;
        cin >> n;
        init();
        int a, b, c;
        repd(i, 1, n)
        {
            head[i] = 0;
            dp[i][0] = dp[i][1] = dp[i][2] = 0;
        }
        repd(i, 2, n)
        {
            cin >> a >> b >> c;
            c %= 3;
            addedge(a, b, c);
            addedge(b, a, c);
        }
        dfs(1, 0);
        cout << ans << endl;
    }


    return 0;
}

inline void getInt(int* p) {
    char ch;
    do {
        ch = getchar();
    } while (ch == ' ' || ch == '\n');
    if (ch == '-') {
        *p = -(getchar() - '0');
        while ((ch = getchar()) >= '0' && ch <= '9') {
            *p = *p * 10 - ch + '0';
        }
    }
    else {
        *p = ch - '0';
        while ((ch = getchar()) >= '0' && ch <= '9') {
            *p = *p * 10 + ch - '0';
        }
    }
}