#include <bits/stdc++.h>
using namespace std;
#define fast ios::sync_with_stdio(0);cin.tie(0)
typedef long long ll;
#define f first
#define s second
#define LOGN 21
const ll MOD = 1e9 + 7;
const ll MAXN = 1e6 + 100;
#define int long long
vector<vector<int>> graph;
map<int,int> mp;
vector<int> A;
int N, v, mx_v, a, b, sz[MAXN], marked[MAXN];
int sum[MAXN], ans = 0;
int get_sz(int node, int parent) {
sz[node] = 1;
for (auto u : graph[node]) {
if (u != parent && !marked[u])
sz[node] += get_sz(u, node);
}
return sz[node];
}
int find_centro(int node, int parent, int n) {
for (auto u : graph[node]) {
if (u != parent && !marked[u] && sz[u] * 2 >= n)
return find_centro(u, node, n);
}
return node;
}
void dfs(int node, int parent, int now, int rem, int len) {
ans = max(ans, now - min(A[rem], A[node]));
if (len == v)
return ;
for (auto u : graph[node]) {
if (u != parent && !marked[u])
dfs(u, node, now + sum[u] - A[node], rem, len + 1);
}
}
void get(int node, int parent, int now, int len) {
for (int i = 2; i <= min(mx_v, v - len); i++)
ans = max(ans, now + mp[i]);
for (auto u : graph[node]) {
if (u != parent && !marked[u])
get(u, node, now + sum[u] - A[node], len + 1);
}
}
void add(int node, int parent, int now, int len) {
mp[len] = max(mp[len], now - A[node]);
mx_v = max(mx_v, len);
for (auto u : graph[node]) {
if (u != parent && !marked[u])
add(u, node, now + sum[u] - A[node], len + 1);
}
}
void decompose(int node) {
int n = get_sz(node, node);
int centro = find_centro(node, node, n);
marked[centro] = true;
dfs(centro, centro, sum[centro] + A[centro], centro, 1);
mp = map<int,int>();
for (int i = 0; i <= 100; i++)
mp[i] = -1e15;
mx_v = 0;
for (auto u : graph[centro]) {
get(u, centro, sum[u] - A[centro], 1);
add(u, centro, sum[centro] + sum[u], 2);
}
mp = map<int,int>();
for (int i = 0; i <= 100; i++)
mp[i] = -1e15;
mx_v = 0;
reverse(graph[centro].begin(), graph[centro].end());
for (auto u : graph[centro]) {
get(u, centro, sum[u] - A[centro], 1);
add(u, centro, sum[centro] + sum[u], 2);
}
}
signed main() {
fast;
cin >> N >> v;
graph = vector<vector<int>>(N+1, vector<int>());
A = vector<int>(N+1);
for (int i = 1; i <= N; i++)
cin >> A[i];
for (int i = 1; i < N; i++) {
cin >> a >> b;
graph[a].push_back(b);
graph[b].push_back(a);
}
for (int i = 1; i <= N; i++) {
for (auto u : graph[i])
sum[i] += A[u];
}
decompose(1);
cout << ans << "\n";
}
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
2396 KB |
Output is correct |
2 |
Incorrect |
1 ms |
2396 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
2396 KB |
Output is correct |
2 |
Incorrect |
1 ms |
2396 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Incorrect |
66 ms |
16856 KB |
Output isn't correct |
2 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
2396 KB |
Output is correct |
2 |
Incorrect |
1 ms |
2396 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |