99 lines
3.1 KiB
C++
99 lines
3.1 KiB
C++
#include <cctype>
|
||
#include <cstdint>
|
||
#include <iostream>
|
||
#include <utility>
|
||
#include <vector>
|
||
#include <map>
|
||
|
||
using ll = int64_t;
|
||
using std::cin, std::cout;
|
||
|
||
// 定义最大节点数为5e5加5，以确保足够的空间
|
||
const ll maxn = 5e5 + 5;
|
||
ll n, k; // n: 树的节点数，k: 暴龙的种类数
|
||
ll a[maxn]; // a[i]: 第i个节点的暴龙种类
|
||
ll ans[maxn]; // ans[i]: 第i种暴龙的红温值
|
||
std::vector<std::pair<ll, ll>> adj[maxn]; // adj[u]: 节点u的邻居及边权
|
||
std::map<ll, ll> dp[maxn]; // dp[u][j]: 节点u到子树中种类j的最大距离
|
||
|
||
// 优化输入读取，提高IO速度
|
||
struct CinNum {
|
||
char c;
|
||
ll n, w;
|
||
// 重载右移运算符以自定义输入方式
|
||
CinNum &operator>>(ll &num) {
|
||
c = n = 0;
|
||
w = 1;
|
||
// 跳过非数字字符，记录负号
|
||
while (!isdigit(c = getchar())) {
|
||
if (c == '-') w = -1;
|
||
}
|
||
// 读取数字
|
||
while (isdigit(c)) {
|
||
n = n * 10 + (c - '0');
|
||
c = getchar();
|
||
}
|
||
num = n * w;
|
||
return *this;
|
||
}
|
||
} cinn;
|
||
// 使用自定义的CinNum来替代标准输入
|
||
#define cin cinn
|
||
|
||
/**
|
||
* @brief 使用深度优先搜索 (DFS) 计算每种暴龙的红温值
|
||
*
|
||
* @param fth 当前节点的父节点
|
||
* @param u 当前处理的节点
|
||
*/
|
||
void dfs(const ll &fth, const ll &u) noexcept {
|
||
// 初始化当前节点的dp，自己到自己距离为0
|
||
dp[u].emplace(a[u], 0);
|
||
|
||
// 遍历所有邻接节点
|
||
for (const auto& [v, w] : adj[u]) {
|
||
if (v == fth) continue; // 避免回到父节点
|
||
dfs(u, v); // 递归处理子节点
|
||
|
||
// 遍历子节点v的所有暴龙种类及其对应的最大距离
|
||
for (const auto& [kk, vv] : dp[v]) {
|
||
// 如果当前节点u已经有这种暴龙类型，且子节点v也有
|
||
if ((dp[u][kk] > 0 || a[u] == kk) && (dp[v][kk] > 0 || a[v] == kk)) {
|
||
// 更新ans[kk]为u和v子树中这种暴龙类型的最大距离
|
||
ans[kk] = std::max(ans[kk], dp[u][kk] + dp[v][kk] + w);
|
||
}
|
||
// 更新u节点到这种暴龙类型的最大距离
|
||
dp[u][kk] = std::max(dp[u][kk], dp[v][kk] + w);
|
||
|
||
// 如果当前节点u本身就是这种暴龙类型，进一步更新ans[kk]
|
||
if (a[u] == kk) {
|
||
ans[kk] = std::max(ans[kk], dp[u][kk]);
|
||
}
|
||
}
|
||
// 清空子节点v的dp以节省内存，因为已经合并到u的dp中了
|
||
dp[v].clear();
|
||
}
|
||
}
|
||
|
||
int main(){
|
||
// 读取节点数n和暴龙种类数k
|
||
cin >> n >> k;
|
||
|
||
// 读取每个节点的暴龙种类
|
||
for(ll i = 1; i <= n; i++) cin >> a[i];
|
||
|
||
// 读取树的边信息，并构建邻接表
|
||
for(ll i = 1; i < n; i++){
|
||
ll u, v, w;
|
||
cin >> u >> v >> w;
|
||
adj[u].emplace_back(v, w);
|
||
adj[v].emplace_back(u, w);
|
||
}
|
||
|
||
// 从节点1开始DFS遍历整棵树，假设节点1为根节点，父节点为0
|
||
dfs(0, 1);
|
||
|
||
// 输出每种暴龙的红温值
|
||
for(ll i = 1; i <= k; i++) cout << ans[i] << '\n';
|
||
}
|