Submission #47211

#	Time	Username	Problem	Language	Result	Execution time	Memory
47211		nickyrio	Chase (CEOI17_chase)	C++17	100 / 100	1693 ms	217036 KiB

chase

#include <bits/stdc++.h>
#define FOR(i, a, b) for (int i = (a); i <= (b); ++i)
#define FORD(i, a, b) for (int i = (a); i >= (b); --i)
#define REP(i, a) for (int i = 0; i < (a); ++i)
#define DEBUG(x) { cerr << #x << '=' << x << endl; }
#define Arr(a, l, r) { cerr << #a << " = {"; FOR(_, l, r) cerr << ' ' << a[_]; cerr << "}\n"; }
#define N 100100
#define pp pair<int, int>
#define endl '\n'
#define IO ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL)
#define taskname ""
#define bit(S, i) (((S) >> (i)) & 1)
void Max(long long &a, long long b) { a = a < b ? b : a; }
void Min(long long &a, long long b) { a = a < b ? a : b; }
using namespace std;

int n, k;
vector<int> e[N];
long long ans, a[N], DtoU[N][111], UtoD[N][111], sum[N];
long long nowUtoD[111], best2UtoD[111], best2DtoU[111], bestUtoD[111], bestDtoU[111];
int skip[N], nChild[N];

void dfs(int u, int p) {
    FOR(i, 1, k) DtoU[u][i] = 0;
    DtoU[u][1] = sum[u];
    for (int v : e[u]) if (v != p && !skip[v]) {
        FOR(i, 1, k) UtoD[v][i] = 0;
        FOR(i, 1, k) {
            // Choose v
            if (i != 1) Max(UtoD[v][i], UtoD[u][i - 1] + sum[v] - a[u]);
            else UtoD[v][i] = sum[v] - a[u];
            // Don't choose v
            Max(UtoD[v][i], UtoD[u][i]);
            Max(nowUtoD[i], UtoD[v][i]);
        }
        dfs(v, u);
        FOR(i, 1, k) {
            // Choose u
            if (i != 1) Max(DtoU[u][i], DtoU[v][i - 1] + sum[u] - a[v]);
            // Don't choose u
            Max(DtoU[u][i], DtoU[v][i]);
        }
    }
}

void dfsCentroid(int u, int p) {
    nChild[u] = 1;
    for (int v : e[u]) if (v != p && !skip[v]) {
        dfsCentroid(v, u);
        nChild[u] += nChild[v];
    }
}

int findCentroid(int u, int p, int total) {
    for (int v : e[u]) if (v != p && !skip[v]) {
        if (nChild[v] * 2 >= total) {
            return findCentroid(v, u, total);
        }
    }
    return u;
}

void Solve(int u) {
    dfsCentroid(u, -1);
    u = findCentroid(u, -1, nChild[u]);
    skip[u] = 1;
    //DEBUG(u);
    FOR(i, 1, k) {
        bestUtoD[i] = bestDtoU[i] = 0;
        best2DtoU[i] = 0;
    }
    Max(ans, sum[u]);
    bool first = true;
    for (int v : e[u]) if (!skip[v]) {
        FOR(i, 0, k) UtoD[v][i] = 0;
        UtoD[v][1] = sum[v] - a[u];
        dfs(v, u);
        // Get ans
        FOR(i, 0, k) {
            // Don't choose u
            Max(ans, bestDtoU[i] + nowUtoD[k - i]);
            Max(ans, bestUtoD[i] + DtoU[v][k - i]);
            // Choose u
            if (i != k) {
                Max(ans, nowUtoD[k - 1 - i] + sum[u]);
                Max(ans, DtoU[v][k - 1 - i] + sum[u] - a[v]);
            }
            if (i && i != k && !first) {
                Max(ans, best2DtoU[i] + sum[u] + nowUtoD[k - 1 - i]);
                Max(ans, bestUtoD[i] + DtoU[v][k - 1 - i] + sum[u] -  a[v]);
            }
        //    cerr << u << ' ' << v  << i << ' ' << ' ' << ans << '\n';
        }
        first = false;
        //cout << u << ' ' << v << ' ' << ans << '\n';
        // Update
        FOR(i, 1, k) {
            Max(bestDtoU[i], DtoU[v][i]);
            Max(bestDtoU[i], bestDtoU[i - 1]);
            Max(bestUtoD[i], nowUtoD[i]);
            Max(bestUtoD[i], bestUtoD[i - 1]);
            Max(best2DtoU[i], DtoU[v][i] - a[v]);
            Max(best2DtoU[i], best2DtoU[i - 1]);
            nowUtoD[i] = 0;
        }
    }
    for (int v : e[u]) if (!skip[v])
        Solve(v);
}
int main() {
    #ifdef NERO
    freopen("test.inp","r",stdin);
    freopen("test.out","w",stdout);
    double stime = clock();
    #else
        //freopen(taskname".inp","r",stdin);
        //freopen(taskname".out","w",stdout);
    #endif //NERO
    IO;
    cin >> n >> k;
    FOR(i, 1, n) cin >> a[i];
    FOR(i, 2, n) {
        int u, v;
        cin >> u >> v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    if (k == 0) {
        cout << 0;
        return 0;
    }
    FOR(u, 1, n) {
        for (int v : e[u]) sum[u] += a[v];
    }
    //Arr(sum, 1, n);
    Solve(1);
    cout << ans;
    #ifdef NERO
    double etime = clock();
    cerr << "\nExecution time: " << (etime - stime) / CLOCKS_PER_SEC * 1000 << " ms.\n";
    #endif // NERO
}

• Subtask 1: 20 / 20
• Subtask 2: 20 / 20
• Subtask 3: 30 / 30
• Subtask 4: 30 / 30