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(2≤n≤1052≤n≤105), the number of cities.

    The following n-1 lines, each contains three integers uiui, vivi, wiwi(1≤ui,vi≤n,1≤wi≤10001≤ui,vi≤n,1≤wi≤1000), the edge of cities.

输出描述:

For each test case print the number of pairs required.

示例1

输入

复制

输出

复制

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 = ; while (b) {if (b % )ans = ans * a % MOD; a = a * a % MOD; b /= ;} return ans;}

inline void getInt(int* p);

const int maxn = ;

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 << ];

ll ans = 0ll;

int tot;

int head[maxn << ];

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 = ;

}

ll dp[maxn][];

void dfs(int id, int pre)

{

    dp[id][] = 1ll;

    for (int i = head[id]; i != ; i = e[i].next)

    {

        edge x = e[i];

        ll w = x.dis;

        if (x.to != pre)

        {

            dfs(x.to, id);

            ans += dp[id][] * dp[x.to][( - w -  + ) % ];

            ans += dp[id][] * dp[x.to][( - w -  + ) % ];

            ans += dp[id][] * dp[x.to][( - w -  + ) % ];

            dp[id][( + w) % ] += dp[x.to][];

            dp[id][( + w) % ] += dp[x.to][];

            dp[id][( + w) % ] += dp[x.to][];

        }

    }

}

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, , n)

        {

            head[i] = ;

            dp[i][] = dp[i][] = dp[i][] = ;

        }

        repd(i, , n)

        {

            cin >> a >> b >> c;

            c %= ;

            addedge(a, b, c);

            addedge(b, a, c);

        }

        dfs(, );

        cout << ans << endl;

    }

    return ;

}

inline void getInt(int* p) {

    char ch;

    do {

        ch = getchar();

    } while (ch == ' ' || ch == '\n');

    if (ch == '-') {

        *p = -(getchar() - '');

        while ((ch = getchar()) >= '' && ch <= '') {

            *p = *p *  - ch + '';

        }

    }

    else {

        *p = ch - '';

        while ((ch = getchar()) >= '' && ch <= '') {

            *p = *p *  + ch - '';

        }

    }

}