update

src/20241118/T541156.cpp

@@ -2,13 +2,33 @@
 #include <iostream>
 #include <utility>
 #include <vector>
+#include <map>
 
 using ll = int64_t;
 using std::cin, std::cout;
 
 const ll maxn = 5e5+5;
-ll n, k, a[maxn];
+ll n, k, a[maxn], ans[maxn];
 std::vector<std::pair<ll, ll>> adj[maxn];
+std::map<ll,ll> dp[maxn];//dp[i][j]:第i个节点距离颜色j的最大值
+
+void dfs(const ll &fth, const ll &u)noexcept{
+    dp[u].emplace(a[u],0);
+    // dp[u][a[u]]=0;
+    for(const auto& [v,w]:adj[u]){
+        if(v==fth)continue;
+        dfs(u,v);
+        for(const auto& [kk,vv]:dp[v]){
+            if((dp[u][kk]>0||a[u]==kk)&&(dp[v][kk]>0||a[v]==kk)){
+                //更新与u,v子树有关的红温值
+                ans[kk]=std::max(ans[kk], dp[u][kk] + dp[v][kk] + w);
+            }
+            dp[u][kk]=std::max(dp[u][kk],dp[v][kk]+w);
+            if(a[u]==kk)ans[kk]=std::max(ans[kk],dp[u][kk]);//更新当前节点的红温值
+        }
+        dp[v].clear();
+    }
+}
 
 int main(){
     std::iostream::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
@@ -21,6 +41,6 @@ int main(){
         adj[u].emplace_back(v,w);
         adj[v].emplace_back(u,w);
     }
-
-    
+    dfs(0,1);
+    for(ll i{1};i<=k;i++)cout<<ans[i]<<'\n';
 }