#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<ll,ll> pl;
typedef pair<int,int> pii;
typedef tuple<int,int,int> tt;
#define all(a) a.begin(), a.end()
#define filter(a) a.erase(unique(all(a)), a.end())
const int mn = 1e5 + 5;
ll pigeon[mn], neighbour[mn], dpDown[2][mn], dpUp[2][mn], dpExcept[2][mn];
vector<int> adj[mn];
struct helper {
vector<ll> pre, sfx;
int n;
helper (vector<ll> a) : pre(a.size()), sfx(a.size()), n(a.size()) {
if (!n) return;
pre[0] = a[0], sfx[n - 1] = a[n - 1];
for (int i = 1; i < n; i++)
pre[i] = max(pre[i - 1], a[i]);
for (int i = n - 2; i >= 0; i--)
sfx[i] = max(sfx[i + 1], a[i]);
}
ll getPre (int p) {
return (0 <= p && p < n ? pre[p] : 0);
}
ll getSfx (int p) {
return (0 <= p && p < n ? sfx[p] : 0);
}
ll query (int exc = -1) {
if (exc == -1) return getPre(n - 1);
return max(getPre(exc - 1), getSfx(exc + 1));
}
};
void dfsDown (int u, int p, int bread) {
int t = bread & 1;
dpDown[t][u] = neighbour[u];
vector<ll> v1(adj[u].size()), v2(adj[u].size());
for (int i = 0; i < adj[u].size(); i++) {
int v = adj[u][i]; if (v == p) continue;
dfsDown(v, u, bread);
v1[i] = dpDown[t][v] - pigeon[u];
v2[i] = dpDown[t ^ 1][v] + neighbour[u] - pigeon[u];
}
helper take(v1), ignore(v2);
dpDown[t][u] = max({dpDown[t][u], take.query(), ignore.query()});
for (int i = 0; i < adj[u].size(); i++) {
int v = adj[u][i]; if (v == p) continue;
dpExcept[t][v] = max({neighbour[u], take.query(i), ignore.query(i)});
}
}
void dfsUp (int u, int p, int bread) {
int t = bread & 1;
dpUp[t][u] = neighbour[u];
if (u != p) {
ll bestTake = max(dpUp[t ^ 1][p], dpExcept[t ^ 1][u]);
ll bestIgnore = max(dpUp[t][p], dpExcept[t][u]);
dpUp[t][u] = max(dpUp[t][u], bestTake + neighbour[u] - pigeon[u]);
dpUp[t][u] = max(dpUp[t][u], bestIgnore - pigeon[u]);
}
for (int v : adj[u])
if (v != p) dfsUp(v, u, bread);
}
int main()
{
ios::sync_with_stdio(0);
cin.tie(0);
int n, v; cin >> n >> v;
for (int i = 1; i <= n; i++) cin >> pigeon[i];
for (int i = 1; i < n; i++) {
int a, b; cin >> a >> b;
neighbour[a] += pigeon[b];
neighbour[b] += pigeon[a];
adj[a].push_back(b);
adj[b].push_back(a);
}
if (v == 0) return cout << 0, 0;
for (int i = 1; i <= v; i++) {
dfsDown(1, 1, i);
dfsUp(1, 1, i);
}
ll ans = 0;
for (int i = 1; i <= n; i++)
ans = max({ans, dpDown[v & 1][i], dpUp[v & 1][i]});
cout << ans;
return 0;
}
Compilation message
chase.cpp: In function 'void dfsDown(int, int, int)':
chase.cpp:49:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
49 | for (int i = 0; i < adj[u].size(); i++) {
| ~~^~~~~~~~~~~~~~~
chase.cpp:58:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
58 | for (int i = 0; i < adj[u].size(); i++) {
| ~~^~~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
2652 KB |
Output is correct |
| 2 |
Incorrect |
1 ms |
2652 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
2652 KB |
Output is correct |
| 2 |
Incorrect |
1 ms |
2652 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
2369 ms |
34132 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
2652 KB |
Output is correct |
| 2 |
Incorrect |
1 ms |
2652 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
