제출 #571884

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
571884		piOOE	Chase (CEOI17_chase)	C++17	100 / 100	648 ms	188892 KiB

chase

#include <bits/stdc++.h>

using namespace std;

#define sz(x) ((int)size(x))
#define all(x) begin(x), end(x)
#define trace(x) cout << #x << ": " << (x) << endl;

typedef long long ll;

mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());

int rand(int l, int r) { return (int) ((ll) rnd() % (r - l + 1)) + l; }

const int N = 100001;
const ll infL = 3e18;

vector<int> g[N];
int n, k;
ll p[N], sum[N], ans = 0;
vector<ll> up[N], down[N];

vector<ll> insert(vector<ll> x, ll v) {
    assert(sz(x) == k + 1);
    for (int i = k; i > 0; --i) {
        x[i] = max(x[i], x[i - 1] + v);
    }
    for (int i = 1; i <= k; ++i) {
        x[i] = max(x[i], x[i - 1]);
    }
    return x;
}

vector<ll> getmax(vector<ll> a, vector<ll> b) {
    assert(sz(a) == k + 1 && sz(b) == k + 1);
    for (int i = 0; i <= k; ++i)
        a[i] = max(a[i], b[i]);
    return a;
}

void dfs(int v, int par) {
    for (int to: g[v]) {
        if (to != par) {
            dfs(to, v);
        }
    }
    for (int cnt = 0; cnt < 2; ++cnt) {
        vector<ll> mx(k + 1);
        for (int to: g[v]) {
            if (to != par) {
                for (int i = 0; i <= k; ++i) {
                    ans = max(ans, down[to][i] + mx[k - i]);
                }
                mx = getmax(mx, insert(up[to], sum[v] - p[to]));
            }
        }
        reverse(all(g[v]));
    }
    up[v] = insert(up[v], sum[v]);
    down[v] = insert(down[v], sum[v] - (par == -1 ? 0 : p[par]));
    for (int to: g[v]) {
        if (to != par) {
            ans = max(ans, up[to][k - 1] + sum[v] - p[to]);
            ans = max(ans, down[to][k - 1] + sum[v]);
            up[v] = getmax(up[v], insert(up[to], sum[v] - p[to]));
            down[v] = getmax(down[v], insert(down[to], sum[v] - (par == -1 ? 0 : p[par])));
        }
    }
    ans = max({ans, up[v][k], down[v][k]});
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> k;
    for (int i = 0; i < n; ++i) {
        cin >> p[i];
        up[i].resize(k + 1);
        down[i].resize(k + 1);
    }
    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        --a, --b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    if (k == 0) {
        cout << 0;
        return 0;
    }
    for (int v = 0; v < n; ++v) {
        for (int to: g[v]) {
            sum[v] += p[to];
        }
    }
    dfs(0, -1);
    cout << ans;
    return 0;
}

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