牛客练习赛30 C 小K的疑惑（01分类+搜索）

Posted 2022-12-11 ITAK

传送门

题目描述

众所周知，小 K是一只连 NOIP2018初赛都没有过的蒟蒻，所以小 K很擅长 dfs序 +分块树，但是本题与 dfs序 +分块树无关。

小K现在心态爆炸了，因为小K被一道简单的数据结构题给卡住了，希望请你来解决它，但是小K又不想太麻烦你，于是将题面进行了简化（其实是出题人懒得写题面了233333）：

Bob有?个点的树，每条边的长度有一个边权，现在定义???(?,?)代表第?个点到第?个点的距离模2之后的结果。问有多少(?,?,?)满足，???(?,?) = ???(?,?) = ???(?,?)。

输入描述:

第一行一个整数?代表点的数量。

接下来? − 1行每行三个数?,?,?代表有一条在?,?之间长度为?的边。

输出描述:

一行一个整数代表有多少对(?,?,?)满足条件。

示例1

输入

复制

3
1 2 3
1 3 4

输出

复制

9

备注:

对于100%的数据，1 ≤ ? ≤ 10000，0 ≤ ? ≤ 233。

解题思路：
因为这是一棵树，所以两个节点之间的路径长度是唯一的，那么我们把 $1$ 号节点当作根节点，求出每个节点到根节点的路径长度（仅用 $01$ 表示即可)，那么就将所有的路径长度分为了两个集合，长度为 $0$ 的集合与长度为 $1$ 的集合，然后我们发现题目要求 $dis(i,j)=dis(j,k)=dis(i,k)$ 的方案数，如果 $dis(i,j)=dis(j,k)$ 那么就有 $dis(i,k)=(dis(i,j)+dis(j,k))\%2==0$ ，所以 $dis(i,j)=dis(j,k)=dis(i,k)=0$，所以我们只需要从 $0$ 和 $1$ 两个集合中取 $3$ 个相同的元素即可，那么方案数 $ans=se{t}_0^3+se{t}_1^3$，我们只需要预处理每个节点到根节点的距离即可。

代码如下：

#include <bits/stdc++.h>
#define pb push_back
#define mk make_pair
#define fi first
#define se second
using namespace std;
typedef pair<int, int> pii;
typedef long long LL;
vector<pii> g[10005];
LL ans;
void dfs(int u, int p, int v) 
    if(v) ans++;
    for(int i=0; i<g[u].size(); i++) 
        if(g[u][i].fi == p) continue;
        dfs(g[u][i].fi, u, v^g[u][i].se);
    

int main() 
    int n; cin>>n;
    for(int i=0; i<n-1; i++) 
        int u, v, w; cin>>u>>v>>w;
        g[u].pb(mk(v, w&1));
        g[v].pb(mk(u, w&1));
    
    ans = 0;
    dfs(1, 0, 0);
    ans = ans*ans*ans+1ll*(n-ans)*(n-ans)*(n-ans);
    cout<<ans<<endl;
    return 0;