This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#define N 100000
#define K 100
long long max(long long a, long long b) { return a > b ? a : b; }
int *ej[N], eo[N];
void append(int i, int j) {
int o = eo[i]++;
if (o >= 2 && (o & o - 1) == 0)
ej[i] = (int *) realloc(ej[i], o * 2 * sizeof *ej[i]);
ej[i][o] = j;
}
int aa[N], k_; long long bb[N], dp[N][K + 1], dq[N][K + 1];
void dfs1(int p, int i) {
int o, k;
memset(dp[i] + 1, -1, k_ * sizeof *dp[i]), dp[i][1] = bb[i];
for (o = eo[i]; o--; ) {
int j = ej[i][o];
if (j != p) {
dfs1(i, j);
for (k = 1; k <= k_; k++)
dp[i][k] = max(dp[i][k], max(dp[j][k], dp[j][k - 1] == -1 ? -1 : dp[j][k - 1] + bb[i] - (k == 1 ? 0 : aa[j])));
}
}
}
void dfs2(int p, int i) {
static long long max1[K + 1], max2[K + 1];
static int jj[K + 1];
int o, k;
if (p == -1)
memset(dq[i] + 1, -1, k_ * sizeof *dq[i]), dq[i][1] = bb[i];
else
for (k = k_; k >= 1; k--)
if (dq[i][k - 1] != -1)
dq[i][k] = max(dq[i][k], dq[i][k - 1] + bb[i] - (k == 1 ? 0 : aa[p]));
memset(max1 + 1, -1, k_ * sizeof *max1);
memset(max2 + 1, -1, k_ * sizeof *max2);
memset(jj + 1, -1, k_ * sizeof *jj);
for (o = eo[i]; o--; ) {
int j = ej[i][o];
if (j != p)
for (k = 1; k <= k_; k++) {
long long x = max(dp[j][k], dp[j][k - 1] == -1 ? -1 : dp[j][k - 1] + bb[i] - (k == 1 ? 0 : aa[j]));
if (max1[k] < x)
max2[k] = max1[k], max1[k] = x, jj[k] = j;
else if (max2[k] < x)
max2[k] = x;
}
}
for (o = eo[i]; o--; ) {
int j = ej[i][o];
if (j != p)
for (k = 1; k <= k_; k++)
dq[j][k] = max(dq[i][k], jj[k] == j ? max2[k] : max1[k]);
}
for (o = eo[i]; o--; ) {
int j = ej[i][o];
if (j != p)
dfs2(i, j);
}
}
int main() {
int n, h, i, j, k;
long long ans;
scanf("%d%d", &n, &k_);
if (k_ == 0) {
printf("0\n");
return 0;
}
for (i = 0; i < n; i++)
ej[i] = (int *) malloc(2 * sizeof *ej[i]);
for (i = 0; i < n; i++)
scanf("%d", &aa[i]);
for (h = 0; h < n - 1; h++) {
scanf("%d%d", &i, &j), i--, j--;
bb[i] += aa[j], bb[j] += aa[i];
append(i, j), append(j, i);
}
dfs1(-1, 0);
dfs2(-1, 0);
ans = 0;
for (i = 0; i < n; i++)
for (k = 0; k <= k_; k++)
ans = max(ans, max(dp[i][k], dq[i][k]));
printf("%lld\n", ans);
return 0;
}
Compilation message (stderr)
chase.c: In function 'append':
chase.c:15:23: warning: suggest parentheses around '-' in operand of '&' [-Wparentheses]
15 | if (o >= 2 && (o & o - 1) == 0)
| ~~^~~
chase.c: In function 'main':
chase.c:83:2: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
83 | scanf("%d%d", &n, &k_);
| ^~~~~~~~~~~~~~~~~~~~~~
chase.c:91:3: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
91 | scanf("%d", &aa[i]);
| ^~~~~~~~~~~~~~~~~~~
chase.c:93:3: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
93 | scanf("%d%d", &i, &j), i--, j--;
| ^~~~~~~~~~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |