#pragma GCC optimize("Ofast,unroll-loops,no-stack-protector")
#include <bits/stdc++.h>
using namespace std;
#define df(b, e) ((b) > (e))
#define fore(i, b, e) for (auto i = (b) - df(b, e); i != e - df(b, e); i += 1 - 2 * df(b, e))
#define sz(x) (int)(x.size())
#define all(x) begin(x), end(x)
#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define mp make_pair
using lli = long long;
using ld = long double;
using ii = pair<int, int>;
using vi = vector<int>;
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...)
#endif
const int N = 1e3 + 5;
const int M = 101;
const int INF = 1e9 + 7;
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
vi g[N];
lli dp[N][M][2], p[N], gans;
bool done[N][M][2];
int k;
template<class T>
void umax(T &a, T b) {
a = max(a, b);
}
// bd -> breadcrumbs
// 1 <- Drop on parent's statue
// 0 <- No Drop on parent's statue
#define state [u][bc][type]
lli dfs (int u, int bc, bool type, int parent = 0) {
auto &ans = dp state;
if (!done state) {
done state = 1;
lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
return tot + p[w];
}) + p[u];
for (auto v: g[u]) {
if (v != parent) {
if (bc) {
umax(ans, res + dfs(v, bc - 1, 1, u) - p[parent] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
umax(ans, dfs(v, bc, 0, u) - (type ? 0 : p[u]));
}
}
}
debug(u, bc, type, ans, res);
}
return ans;
}
lli f (int u, int croot, int bc, bool type) {
lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
return tot + p[w];
}) + p[u];
lli ans = 0;
for (int v: g[u]) {
if (v != croot) {
if (bc) {
umax(ans, res + dfs(v, bc - 1, 1, u) - p[croot] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
umax(ans, dfs(v, bc, 0, u) - (type ? 0 : p[u]));
}
}
}
dp state = ans;
done state = 1;
return ans;
}
void solve (int u, int parent = 0) {
lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
return tot + p[w];
}) + p[u];
lli ans = 0;
for (int v: g[u]) {
if (v != parent) {
// umax(ans, res + dfs(v, k - 1, 1, u) - p[parent] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
// dp[u][k][0] = max(res + dfs(v, k - 1, 1, u) - p[u];
umax(ans, res + dfs(v, k - 1, 1, u) - p[u]);
umax(ans, dfs(v, k, 0, u) - p[u]);
} else {
umax(ans, res + f(parent, u, k - 1, 1) - p[u]);
umax(ans, f(parent, u, k, 0) - p[u]);
}
}
debug(u, ans, dp[u][k][1]);
dp[u][k][0] = ans;
umax(gans, ans);
for (int v: g[u]) {
if (v != parent) {
solve(v, u);
}
}
}
int main() {
cin.tie(0) -> sync_with_stdio(0), cout.tie(0);
#ifdef LOCAL
auto start_time = clock();
#endif
int n, u, v;
cin >> n >> k;
fore (i, 1, n + 1)
cin >> p[i];
fore (i, 1, n) {
cin >> u >> v;
g[u].eb(v);
g[v].eb(u);
}
if (k == 0) {
cout << 0 << '\n';
return 0;
}
//cout << dfs(6, k, 0, 0) << '\n';
fore (u, 1, n + 1) {
memset(dp, 0, sizeof(dp));
memset(done, 0, sizeof(done));
umax(gans, dfs(u, k, 0, 0));
}
cout << gans << '\n';
#ifdef LOCAL
auto end_time = clock();
cerr << "Execution time: " << (end_time - start_time) * (int)1e3 / CLOCKS_PER_SEC << " ms\n";
#endif
return 0;
}
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
2176 KB |
Output is correct |
2 |
Correct |
2 ms |
2176 KB |
Output is correct |
3 |
Correct |
2 ms |
2176 KB |
Output is correct |
4 |
Correct |
2 ms |
2176 KB |
Output is correct |
5 |
Correct |
1 ms |
384 KB |
Output is correct |
6 |
Correct |
2 ms |
2176 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
2176 KB |
Output is correct |
2 |
Correct |
2 ms |
2176 KB |
Output is correct |
3 |
Correct |
2 ms |
2176 KB |
Output is correct |
4 |
Correct |
2 ms |
2176 KB |
Output is correct |
5 |
Correct |
1 ms |
384 KB |
Output is correct |
6 |
Correct |
2 ms |
2176 KB |
Output is correct |
7 |
Execution timed out |
4090 ms |
2304 KB |
Time limit exceeded |
8 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
13 ms |
2176 KB |
Execution killed with signal 11 |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
2 ms |
2176 KB |
Output is correct |
2 |
Correct |
2 ms |
2176 KB |
Output is correct |
3 |
Correct |
2 ms |
2176 KB |
Output is correct |
4 |
Correct |
2 ms |
2176 KB |
Output is correct |
5 |
Correct |
1 ms |
384 KB |
Output is correct |
6 |
Correct |
2 ms |
2176 KB |
Output is correct |
7 |
Execution timed out |
4090 ms |
2304 KB |
Time limit exceeded |
8 |
Halted |
0 ms |
0 KB |
- |