99 lines
2.0 KiB
C++
99 lines
2.0 KiB
C++
#include <cstdint>
|
|
#include <iostream>
|
|
#include <vector>
|
|
#include <bitset>
|
|
|
|
using ll = int64_t;
|
|
using std::cin, std::cout;
|
|
|
|
const ll maxn=5e5+5, L{20};
|
|
ll n, k, a[maxn], fa[maxn][L], dpt[maxn], w[maxn];
|
|
std::vector<ll> g[maxn];
|
|
std::vector<std::pair<ll,ll>> adj[maxn];
|
|
std::bitset<maxn> vis;
|
|
|
|
void dfs(const ll &fth,const ll &u){
|
|
fa[u][0]=fth;
|
|
for(auto& [v,ww]:adj[u]){
|
|
if(v==fth)continue;
|
|
dpt[v]=dpt[u]+1;
|
|
w[v]=w[u]+ww;
|
|
dfs(u,v);
|
|
}
|
|
}
|
|
|
|
void initfa(){
|
|
for(ll k{1};k<L;k++){
|
|
for(ll u{1};u<=n;u++){
|
|
fa[u][k]=fa[fa[u][k-1]][k-1];
|
|
}
|
|
}
|
|
}
|
|
|
|
ll lca(ll u,ll v){
|
|
if(dpt[u]<dpt[v]){
|
|
std::swap(u,v);
|
|
}
|
|
ll dif = dpt[u]-dpt[v];
|
|
for(ll j{0};dif;dif>>=1,j++){
|
|
if(dif&1){
|
|
u=fa[u][j];
|
|
}
|
|
}
|
|
if(u==v){
|
|
return u;
|
|
}
|
|
for(ll j{L-1};j>=0;j--){
|
|
if(fa[u][j]!=fa[v][j]){
|
|
u=fa[u][j];
|
|
v=fa[v][j];
|
|
}
|
|
}
|
|
return fa[u][0];
|
|
}
|
|
|
|
|
|
ll dis(const ll &u,const ll &v){
|
|
return w[u]+w[v]-2*w[lca(u,v)];
|
|
}
|
|
|
|
int main(){
|
|
std::iostream::sync_with_stdio(false), std::cin.tie(nullptr), std::cout.tie(nullptr);
|
|
|
|
cin>>n>>k;
|
|
for(ll i{1};i<=n;i++){
|
|
cin>>a[i];
|
|
g[a[i]].push_back(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);
|
|
}
|
|
dfs(0,1);
|
|
initfa();
|
|
for(ll i{1};i<=k;i++){
|
|
if(g[i].size()<2){
|
|
cout<<"0\n";
|
|
continue;
|
|
}
|
|
ll u{g[i][0]},v{g[i][1]};
|
|
ll ans{dis(u,v)};
|
|
for(ll j{2};j<(ll)g[i].size();j++){
|
|
ll du{dis(u,g[i][j])},dv{dis(v,g[i][j])};
|
|
if(du>dv){
|
|
if(du>ans){
|
|
v=g[i][j];
|
|
ans=du;
|
|
}
|
|
}else{
|
|
if(dv>ans){
|
|
ans=dv;
|
|
u=g[i][j];
|
|
}
|
|
}
|
|
}
|
|
cout<<ans<<'\n';
|
|
}
|
|
} |