제출 #571771

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
571771		piOOE	Chase (CEOI17_chase)	C++17	70 / 100	92 ms	12552 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], ans = 0;

multiset<ll> st;

ll sum_set = 0;

void dfs(int v, int par) {
    ll sum = 0;
    for (int to : g[v]) {
        if (to != par) {
            sum += p[to];
        }
    }
    bool add = false;
    ll del = -1;
    if (sz(st) < k) {
        st.insert(sum);
        sum_set += sum;
        add = true;
    } else {
        assert(!st.empty());
        if (*st.begin() < sum) {
            del = *st.begin();
            sum_set -= del;
            st.erase(st.begin());
            add = true;
            sum_set += sum;
            st.insert(sum);
        }
    }
    ans = max(ans, sum_set);
    for (int to: g[v]) {
        if (to != par) {
            dfs(to, v);
        }
    }
    if (add) {
        sum_set -= sum;
        st.extract(sum);
    }
    if (del != -1) {
        sum_set += del;
        st.insert(del);
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> k;
    for (int i = 0; i < n; ++i) {
        cin >> p[i];
    }
    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;
    }
    if (n <= 1000) {
        for (int s = 0; s < n; ++s) {
            dfs(s, -1);
        }
        cout << ans;
    } else {
        dfs(0, -1);
        cout << ans;
    }
    return 0;
}

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