牛客 216 C 小K的疑惑

大意: 给定树, 求多少个三元组(i,j,k), 满足dis(i,j)=dis(j,k)=dis(k,i).

刚开始想复杂了, 暴力统计了所有的情况.

#include <iostream>
#include <queue>
#include <cstdio>
#define REP(i,a,n) for(int i=a;i<=n;++i)
using namespace std;
typedef long long ll;

const int N = 1e4+10;
int n, dep[N];
struct _ {int to,w;};
vector<_> g[N];
int dp[N][2];
ll ans;

void dfs(int x, int fa, int d) {
    dep[x] = d, ++ans;
    ll c00 = 0, c11 = 0;
    for (_ e:g[x]) if (e.to!=fa) {
        int y = e.to;
        dfs(y,x,d+e.w&1);
        ans += 6ll*c00*dp[y][0];
        ans += 6ll*c11*dp[y][1];
        ans += 6*dp[y][dep[x]];
        ans += 12ll*dp[x][dep[x]]*dp[y][dep[x]];
        ans += 6ll*dp[y][dep[x]]*(dp[y][dep[x]]-1)/2;
        ans += 6ll*dp[x][!dep[x]]*dp[y][!dep[x]];
        c00 += (ll)dp[y][0]*dp[x][0];
        c11 += (ll)dp[y][1]*dp[x][1];
        dp[x][0] += dp[y][0];
        dp[x][1] += dp[y][1];
    }
    for (_ e:g[x]) if (e.to!=fa) {
        int y = e.to;
        ans += 6ll*dp[y][0]*(dp[y][0]-1)/2*(dp[x][0]-dp[y][0]);
        ans += 6ll*dp[y][1]*(dp[y][1]-1)/2*(dp[x][1]-dp[y][1]);
    }
    ++dp[x][dep[x]];
}

int main() {
    scanf("%d", &n);
    REP(i,2,n) {
        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,0);
    printf("%lld\n", ans);
}

实际上可以发现所有路径都满足 奇奇奇 或 偶偶偶.

#include <iostream>
#include <queue>
#include <cstdio>
#define REP(i,a,n) for(int i=a;i<=n;++i)
using namespace std;
typedef long long ll;

const int N = 1e4+10;
int n, dep[N];
struct _ {int to,w;};
vector<_> g[N];
int dp[N][2], ans[2];

void dfs(int x, int fa, int d) {
	++ans[d];
    for (_ e:g[x]) if (e.to!=fa) {
        int y = e.to;
        dfs(y,x,d+e.w&1);
    }
}

int main() {
    scanf("%d", &n);
    REP(i,2,n) {
        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,0);
	printf("%lld\n",(ll)ans[0]*ans[0]*ans[0]+(ll)ans[1]*ans[1]*ans[1]);
}