#include <bits/stdc++.h>
#define int long long
using namespace std;
const int nmax = 1e5 + 5, vmax = 1e2 + 5, kmax = 1e2 + 5;
vector<int> g[nmax];
int p[nmax], nval[nmax], dval[nmax], mx[nmax][vmax], altmx[nmax][vmax], atrmx[nmax][vmax];
int main() {
int n, k;
cin >> n >> k;
for(int i = 1; i <= n; i++)
cin >> nval[i];
for(int i = 0, x, y; i < n; i++) {
cin >> x >> y;
g[x].push_back(y);
g[y].push_back(x);
}
for(int i = 1; i <= n; i++) {
for(auto x : g[i])
dval[i] += nval[x];
}
auto initdp = [&](auto &&self, int node, int f) -> void {
mx[node][0] = dval[node];
altmx[node][0] = 0;
atrmx[node][0] = node;
p[node] = f;
for(int i = 0; i < g[node].size(); i++) {
if(g[node][i] == f)
swap(g[node][i], g[node].back()),
g[node].pop_back(),
i--;
else
self(self, g[node][i], node);
}
for(int i = 1; i <= k; i++) {
mx[node][i] = mx[node][i - 1];
altmx[node][i] = altmx[node][i - 1];
atrmx[node][i] = atrmx[node][i - 1];
}
for(auto x : g[node]) {
for(int i = 1; i <= k; i++) {
int curr = mx[x][i - 1] + dval[node];
if(curr >= mx[node][i]) {
altmx[node][i] = mx[node][i];
atrmx[node][i] = x;
mx[node][i] = curr;
}
else if(curr > altmx[node][i])
altmx[node][i] = curr;
}
}
//cerr << node << ":\n";
//for(int i = 0; i <= k; i++)
//cerr << mx[node][i] << ' ' << altmx[node][i] << '\n';
//cerr << '\n';
return;
};
initdp(initdp, 1, 1);
int totalmx = 0;
auto findval = [&](auto &&self, int node) -> void {
int collat = dval[node], last = node, f = p[node];
totalmx = max(totalmx, mx[node][k - 1] - nval[node]);
for(int i = 1; i < k; i++) {
if(last == f)
break;
//if(node == 6) {
//cout << atrmx[f][k - i - 1] << ' ' << last << ' ' << f << '\n';
//cout << mx[f][k - i - 1] << ' ';
//}
totalmx = max(totalmx, collat + (last == atrmx[f][k - i - 1]? altmx : mx)[f][k - i - 1] - nval[node]);
collat += dval[f];
last = p[last];
f = p[f];
}
for(auto x : g[node])
self(self, x);
return;
};
findval(findval, 1);
cout << totalmx << '\n';
}
Compilation message
cc1plus: error: '::main' must return 'int'
chase.cpp: In instantiation of 'main()::<lambda(auto:1&&, long long int, long long int)> [with auto:1 = main()::<lambda(auto:1&&, long long int, long long int)>&]':
chase.cpp:63:22: required from here
chase.cpp:31:22: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
31 | for(int i = 0; i < g[node].size(); i++) {
| ~~^~~~~~~~~~~~~~~~
