제출 #468646

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
468646		idk321	Chase (CEOI17_chase)	C++17	40 / 100	941 ms	524292 KiB

chase

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N = 100005;
const int V = 101;
int p[N];
int n, v;
vector<int> adj[N];


struct comp {
    bool operator () (const array<ll,2>& ar1, const array<ll,2>& ar2) const {
        return ar1 > ar2;
    }
};

bool comp_fn (const array<ll,2>& ar1, const array<ll,2>& ar2)  {
    return ar1 > ar2;
}


struct SetLike {
    vector<array<ll, 2>> v;

};


vector<array<ll, 2>>::iterator sbegin(vector<array<ll, 2>>& v) {
        return v.begin();
}

vector<array<ll, 2>>::iterator send(vector<array<ll, 2>>& v) {
    return v.end();
}


void sinsert(vector<array<ll, 2>>& v, array<ll, 2> ar) {
    v.push_back(ar);
    sort(v.begin(), v.end(), comp_fn);
    if (v.size() > 2) v.pop_back();
}

vector<array<ll, 2>> dpSet[N][V];
ll dp[N][V];

ll defaultValue[N];

void dfs(int node, int par) {
    ll val1 = 0;
    for (int next : adj[node]) {
        defaultValue[node] += p[next];
        if (next == par) {
            continue;
        }
        val1 += p[next];
        dfs(next, node);
    }

    dp[node][1] = max(dp[node][1], val1);
    for (int next : adj[node]) {
        if (next == par) continue;
        for (int i = 0; i <= v; i++) {
            sinsert(dpSet[node][i], {dp[next][i], next});
            if (i != v) dp[node][i + 1] = max(dp[node][i + 1], dp[next][i] + val1);
            dp[node][i] = max(dp[node][i], dp[next][i]);
        }
    }

    for (int i = 0; i <= v; i++) {
        sinsert(dpSet[node][i], {0, node});
        sinsert(dpSet[node][i], {0, node});
    }
}

void transferRoot(int first, int second) {

    for (int i = 0; i <= v; i++) {
        auto it0 = sbegin(dpSet[first][i]);
        if ((*it0)[1] == second) it0++;
        ll val0 = 0;
        if (it0 != send(dpSet[first][i])) val0 = (*it0)[0];
        sinsert(dpSet[second][i], {val0, first});
        dp[second][i] = max(dp[second][i], val0);
        if (i != v)  {
            sinsert(dpSet[second][i + 1], {val0 + defaultValue[first] - p[second], first});
            dp[second][i + 1] = max(dp[second][i + 1], val0 + defaultValue[first] - p[second]);
            dp[second][i + 1] = max(dp[second][i + 1], val0 + defaultValue[second]);
        }
        if (i + 2 <= v) {
            dp[second][i + 2] = max(dp[second][i + 2], val0 + defaultValue[first] - p[second] + defaultValue[second]);
        }
    }
    for (int next : adj[second]) {
        if (next == first) continue;
        for (int i = 0; i <= v; i++) {

            if (i != v) dp[second][i + 1] = max(dp[second][i + 1], dp[next][i] + defaultValue[second]);
            dp[second][i] = max(dp[second][i], dp[next][i]);
        }
    }
}

ll result = 0;

void dfsRoot(int node, int par) {
    for (int i = 0; i <= v; i++) {

        result = max(result, dp[node][i]);

    }


    for (int next : adj[node]) {
        if (next == par) continue;

        transferRoot(node, next);
        dfsRoot(next, node);
    }
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> v;
    for (int i = 1; i <= n; i++) cin >> p[i];
    for (int i = 0; i < n - 1; i++) {
        int a, b;
        cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    dfs(1, -1);

    dfsRoot(1, -1);
    cout << result << "\n";
}

/*
12 2
2 3 3 8 1 5 6 7 8 3 5 4
2 1
2 7
3 4
4 7
7 6
5 6
6 8
6 9
7 10
10 11
10 12
*/

/*
10 3
10 1 1 1 1 10 1 10 1 10
1 2
2 3
3 4
4 5
5 6
3 7
7 8
9 4
9 10
*/

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