牛客算法周周练18A - 小K的疑惑（Dijkstra + 01分类）

题目大意：

给出一个 n 个节点的树，输入n，再输入n - 1条边，每条边输入u v w表示u到v之间有一条权值为w的边，你需要找到三元组（i, j, k）使得dis(i, j) = dis(i, k) = dis(j, k) 。

解题思路：

观察题意，两点的距离是需要 % 2的，所以说距离非0即1，而i j k 是可以重复的，我们分为0 1 两类：即一个点到节点1的距离的奇偶性。拿奇数为例，比如有n个小球，需要放入（i, j, k）这样的三元组，每个小球可以重复使用，所以一共有n³种方案，偶数也是如此，设到起点的距离为奇数的为ans1, 偶数为ans2, 则最终答案就是ans1³ + ans2³。

Code：

#pragma GCC optimize(2)
#include <cstdio>
#include <iostream>
#include <algorithm>
#include <cstring>
#include <sstream>
#include <cmath>
#include <vector>
#include <queue>
#include <stack>
#include <map>
#include <set>
#include <unordered_map>
#include <unordered_set>
#define pm make_pair
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
const int N = 1e4 + 50;
const int mod = 1e9 + 7;
const int inf = 0x3f3f3f3f;
struct node
{
	int to, dir;
	bool operator < (const node &t) const
	{
		return dir > t.dir;
	}
};
vector<node > e[N];
int n, dis[N];
bool vis[N];

void dijkstra(int s)
{
	for (int i = 1; i <= n; i ++)
		dis[i] = (i == s ? 0 : inf);
	memset(vis, false, sizeof vis);
	priority_queue<node > q;
	q.push({s, 0});
	while (!q.empty())
	{
		auto t = q.top();
		q.pop();
		if (vis[t.to])  continue;
		vis[t.to] = true;
		for (int i = 0; i < e[t.to].size(); i ++)
		{
			int k = e[t.to][i].to;
			if (e[t.to][i].dir + t.dir < dis[k])
			{
				dis[k] = e[t.to][i].dir + t.dir;
				q.push({k, dis[k]});
			}
		}
	}
}

int main()
{
	ios::sync_with_stdio(false);
	cin >> n;
	for (int i = 1; i <= n; i ++)
	{
		int u, v, w;
		cin >> u >> v >> w;
		e[u].push_back({v, w});
		e[v].push_back({u, w});
	}
	dijkstra(1);
	ll ans1 = 0, ans2 = 0;
	for (int i = 1; i <= n; i ++)
		dis[i] % 2 ? ans1++ : ans2++;
	ll ans3 = pow(ans1, 3) + pow(ans2, 3);
	cout << ans3 << endl;
	return 0;
}