did this question come in any OA?

I was thinking of solving it using DP.For each node, I keep a dp[u][k], which means the maximum gain we can get from the subtree of u if we do exactly k operations inside it. It's like a knapsack: each child subtree gives us some set of gains based on how many moves we spend on it,and we want to combine them to get the best possible total using at most K moves. O(n*k*(work on each state))= o(n*k*(nk*k))

Gain : let's say u have a leaf at depth d when u reattach it to root overall score will decrease by (d-1)*node_val We can precompute this using dfs in o(n)

If some one has better approsch PLZZ SHARE

try to minimize the contribution of each node to the final answer

and contribution will be depth*a[i] or (depth+1)*a[i] according to your depth definiton

so remove the leaf with max value of depth*a[i] and put it to node ..

repeat the operation k times..

Are the values at the nodes all positive ? And what about the constraints ?

Quite Nice Problem ! Here is my solution O(N*K) EULER TOUR + KNAPSACK

It is obvious that the contribution of a particular node is (depth[i] + 1) * val[i] (Considering depth of root to be 0)

Now some observations : ->For every node if at all we choose to apply operation then most optimal will be to connect that node directly to the root [Since that way contribution becomes 2 * val[i] and that's the min possible] ->If we apply the operation on a node then have to apply on every node in its subtree

So compute the gain (depth[i] — 1) * val[i] for all nodes And then sum all the gains from a subtree on the subtree's root also compute the subtree sizes

Let that array formed be arr Now perform EULER TOUR on the tree and let that array be 'e' (All subtree nodes follow their root node) now make the array a where a[i] = arr[e[i]]

Now apply knapsack on a and dp[i][k] = max(dp[i+1][k], a[i] + dp[i+ss[i]][k-ss[i]) ss[i] = subtree size of 'i'

Do you remember the other problems ?

This solution would work I think. Used priority queue and then selected the top k contributors leaving beside the root and connected them to the root so that their contribution minimizes. contribution is just depth[node]*value[node]

void solve() {
  int n, k;
  cin >> n >> k;
  vector<int> a(n + 1);
  for (int i = 1; i <= n; i++) cin >> a[i];
  vector<vector<int>> g(n + 1);
  vector<int> deg(n + 1);
  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
    deg[u]++;
    deg[v]++;
  }

  vector<int> depth(n + 1), par(n + 1);
  queue<int> q;
  q.push(1);
  depth[1] = 1;
  while (!q.empty()) {
    int u = q.front();
    q.pop();
    for (int v : g[u]) {
      if (depth[v] == 0) {
        depth[v] = depth[u] + 1;
        par[v] = u;
        q.push(v);
      }
    }
  }

  long long ans = 0;
  for (int i = 1; i <= n; i++) ans += 1LL * a[i] * depth[i];

  priority_queue<pair<long long, int>> pq;
  for (int i = 2; i <= n; i++) {
    if (deg[i] == 1) {
      long long b = 1LL * a[i] * (depth[i] - 2);
      if (b > 0) pq.push({b, i});
    }
  }

  while (k-- and !pq.empty()) {
    auto [b, u] = pq.top();
    pq.pop();
    if (b <= 0) break;
    ans -= b;
    int p = par[u];
    deg[p]--;
    if (p != 1 and deg[p] == 1) {
      long long b2 = 1LL * a[p] * (depth[p] - 2);
      if (b2 > 0) pq.push({b2, p});
    }
  }

  cout << ans;
}

I think this works fine. If some mistake is there please rectify.

#include<bits/stdc++.h>
#define cb(x) __builtin_popcountll(x)
#define ctz(x) __builtin_ctzll(x)
#define clz(x) __builtin_clzll(x)
#define int long long
#define ll long long
#define pb push_back
#define debug(x) cout<<"Debug: "<<x<<endl;
using namespace std;

int mpow(int a,int b,int m){int x=1;while(b>0){if(b&1){x=((ll)x*a)%m;}a=((ll)a*a)%m;b>>=1;}return x;}
int mdiv(int a,int b, int m){int bin=mpow(b,m-2,m);return ((ll)a*bin)%m;}
int fact(int n,int m){if(n<=0) return 1;return ((ll)n*fact(n-1,m))%m;}
int nCr(int n,int r,int m){if(r==0)return 1;int ans=fact(n,m);ans=mdiv(ans,fact(n-r,m),m);ans=mdiv(ans,fact(r,m),m);return ans;}
int GCD(int a,int b){if(b==0)return a;return GCD(b,a%b);}

#define INF 1e17
vector<vector<int>> adj;
vector<int> val;
vector<vector<int>> dp;
vector<int> sz;

void dfs(int node,int par,int depth,int k){
   if(adj[node].size() == 1 && par!=-1){
      sz[node] = 1;
      dp[node][0] = 0;
      dp[node][1] = val[node]*(depth-1);
      return;
   }

   for(auto &ch:adj[node]){
      if(ch == par)continue;
      dfs(ch,node,depth+1,k);
      sz[node]+=sz[ch]; 
   }
   sz[node]++;

   for(auto &ch:adj[node]){
      if(ch == par)continue;
      for(int i=0;i<=k;i++){
         for(int j=0;j<=i;j++){
            if(sz[node]==i){
               dp[node][i] = max({dp[node][i], dp[node][j] + dp[ch][i-j], (val[node]*(depth-1)) + dp[ch][i-1]});
            }
            else dp[node][i] = max({dp[node][i], j>0?dp[node][j]:0 + dp[ch][i-j]});
         }
      }
   }
}
void solve(){
   int n,k;
   cin>>n>>k;
   adj.assign(n+1,{}); val.assign(n+1,0);
   sz.assign(n+1,0);
   for(int i=0;i<n-1;i++){
      int u,v;
      cin>>u>>v;
      adj[u].pb(v);
      adj[v].pb(u);
   }
   for(int i=1;i<=n;i++){
      cin>>val[i];
   }
   dp.assign(n+1,vector<int>(k+1,-INF));
   dfs(1,-1,0,k);
   cout<<dp[1][k]<<endl;
}
int32_t main() {
   ios_base::sync_with_stdio(false);
   cin.tie(NULL);
       solve();
    return 0;
}