Submission #873896

#	Time	Username	Problem	Language	Result	Execution time	Memory
873896		noiaint	Chase (CEOI17_chase)	C++17	100 / 100	266 ms	252680 KiB

chase

#include <bits/stdc++.h>
#define int long long
 
using namespace std;
 
#define file "chase"
 
#define mp make_pair
#define fi first
#define se second
#define all(x) x.begin(), x.end()
 
#define getbit(x, i) (((x) >> (i)) & 1)
#define bit(x) (1LL << (x))
#define popcount __builtin_popcountll
 
mt19937_64 rd(chrono::steady_clock::now().time_since_epoch().count());
int rand(int l, int r) {
    return l + rd() % (r - l + 1);
}
 
const int N = 1e5 + 5;
const int mod = (int)1e9 + 7; // 998244353;
const int lg = 25; // lg + 1
const int oo = 1e9;
const long long ooo = 1e18;
 
template<class X, class Y> bool mini(X &a, Y b) {
    return a > b ? (a = b, true) : false;
}
template<class X, class Y> bool maxi(X &a, Y b) {
    return a < b ? (a = b, true) : false;
}
void add(int &a, int b) {
    a += b;
    if (a >= mod) a -= mod;
    if (a < 0) a += mod;
}
 
int n, k;
vector<int> adj[N];
int a[N];
int f[N][105][2];
/*
f[u][i][0] : best value
f[u][i][1] : second-best value for rerooting
*/
int sum[N];
int tmp[N][105];
 
int res = 0;
 
void dfsF(int u, int p) {
    for (int v : adj[u]) if (v != p) {
        dfsF(v, u);
        for (int i = 1; i <= k; ++i) {
            int d = max(f[v][i][0], f[v][i - 1][0] + sum[v] - a[u]);
            if (d > f[u][i][0]) f[u][i][1] = f[u][i][0], f[u][i][0] = d;
            else maxi(f[u][i][1], d);
        }
    }
}
 
void dfsG(int u, int p) {
    for (int i = 1; i <= k; ++i) {
        maxi(res, f[u][i][0]);
        maxi(res, f[u][i - 1][0] + sum[u]);
    }
    for (int v : adj[u]) if (v != p) {
        // save changes
        for (int i = 1; i <= k; ++i) tmp[u][i] = f[u][i][0];
 
        // remove contribution
        for (int i = 1; i <= k; ++i) {
            int d = max(f[v][i][0], f[v][i - 1][0] + sum[v] - a[u]);
            if (f[u][i][0] == d) f[u][i][0] = f[u][i][1];            
        }
 
        // reroot
        for (int i = 1; i <= k; ++i) {
            int d = max(f[u][i][0], f[u][i - 1][0] + sum[u] - a[v]);
            if (d > f[v][i][0]) f[v][i][1] = f[v][i][0], f[v][i][0] = d;
            else maxi(f[v][i][1], d);
        }
        dfsG(v, u);
 
        // rollback changes
        for (int i = 1; i <= k; ++i) f[u][i][0] = tmp[u][i];
 
    }
}
 
signed main() {
    ios::sync_with_stdio(false);
    cin.tie(0);
 
    // freopen(file".inp", "r", stdin);
    // freopen(file".out", "w", stdout);
 
    cin >> n >> k;
    for (int i = 1; i <= n; ++i) cin >> a[i];
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
 
    for (int u = 1; u <= n; ++u) {
        for (int v : adj[u]) sum[u] += a[v];
    }
    
    dfsF(1, 0);
    dfsG(1, 0);
 
    cout << res;
 
    return 0;
}
 
/*
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
 
36
*/

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