小K的疑惑 - 牛客

题意

Bob有\(N(1 \leq N \leq10000)\)个点的树，每条边有一个边权\(d(0\leq d \leq 233)\)，现在定义\(dis(i,j)\)代表第\(i\)个点到第\(j\)个点的距离模2。问有多少\((i,j,k)\)满足\(dis(i,j) = dis(i,k) = dis(j,k)\)。

题解

首先，一棵树中不存在\((i,j,k)\)，使得\(dis(i,j) = dis(i,k) = dis(j,k) = 1\)。随便选一个点为根，记\(dis(root, u) = 0\)的个数为\(cnt_x\)，\(dis(root, u) = 1\)的个数为\(cnt_y\)，那么满足\(dis(i,j) = dis(i,k) = dis(j,k) = 0\)的\((i, j,k)\)个数为\(cnt_x^3 + cnt_y^3\)。注意\(i = j = k\)是合法的！

代码

const int N = 10005;

int n, x, y;

vector<P> G[N];

void DFS(int u, int p, int deep) {
    (deep & 1) ? ++x : ++y;
    for (auto v : G[u]) if (v.first != p) DFS(v.first, u, deep + v.second);
}

void addedge(int u, int v, int w) {
    G[u].pb(P(v, w));
    G[v].pb(P(u, w));
}

int main()
{
    cin >> n;
    rep(i, 1, n) {
        int u, v, w;
        cin >> u >> v >> w;
        addedge(u, v, w);
    }
    DFS(1, 1, 0);
    cout << 1ll * x * x * x + 1ll * y * y * y << endl;
    return 0;
}