Submission #143306

#	Time	Username	Problem	Language	Result	Execution time	Memory
143306		Milki	Chase (CEOI17_chase)	C++14	100 / 100	1228 ms	335268 KiB

chase

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

#define FOR(i, a, b) for(int i = a; i < b; ++i)
#define REP(i, n) FOR(i, 0, n)
#define _ << " " <<
#define sz(x) ((int) x.size())
#define pb(x) push_back(x)
#define TRACE(x) cerr << #x << " = " << x << endl

typedef long long ll;
typedef pair<ll, ll> point;

const int MAXN = 1e5 + 5, MAXV = 105;
const ll inf = 1e16;

int n, k, val[MAXN];
vector <int> E[MAXN];
ll dp_down[MAXV][MAXN], sum[MAXN], sol, dp_up[MAXV][MAXN];

struct Best{
  point a, b;
  Best(){
    a = point(0, MAXN - 1);
    b = point(0, MAXN - 1);
  }
  void ubaci(point x){
    if(x.first > a.first){
      if(x.second != a.second)
        b = a;

      a = x;
    }
    else if(x.first > b.first){
      if(x.second != a.second)
        b = x;
    }
  }
};

void dfs(int x, int p = -1){
  vector <Best> mx1, mx2;
  mx1.resize(k + 1); mx2.resize(k + 1);

  dp_up[1][x] = sum[x];
  mx1[1].ubaci( {sum[x], -1} );

  for(auto e : E[x]){
    if(e == p) continue;

    dfs(e, x);
    FOR(i, 1, k + 1){
      ll ans = dp_down[i - 1][e] + sum[e] - val[x];
      ans = max(ans, dp_down[i][e]);

      if(ans > dp_down[i][x]){
        dp_down[i][x] = ans;
      }
      mx2[i].ubaci(point(ans, e));

      dp_down[i][x] = max(dp_down[i][x], dp_down[i - 1][x]);

      ans = dp_up[i - 1][e] + sum[x] - val[e];
      ans = max(ans, dp_up[i][e]);

      if(ans > dp_up[i][x]){
        dp_up[i][x] = ans;
      }
      mx1[i].ubaci(point(ans, e));

      dp_up[i][x] = max(dp_up[i][x], dp_up[i - 1][x]);

      mx1[i].ubaci(mx1[i - 1].a);
      mx1[i].ubaci(mx1[i - 1].b);

      mx2[i].ubaci(mx2[i - 1].a);
      mx2[i].ubaci(mx2[i - 1].b);
    }
  }

  REP(i, k + 1){
    if(mx1[i].a.second == mx2[k - i].a.second){
      sol = max(sol, mx1[i].a.first + mx2[k - i].b.first);
      sol = max(sol, mx1[i].b.first + mx2[k - i].a.first);
    }
    else{
      sol = max(sol, mx1[i].a.first + mx2[k - i].a.first);
    }
  }

}

signed main(){
  ios_base::sync_with_stdio(false); cin.tie(0);

  cin >> n >> k;
  REP(i, n)
    cin >> val[i];
  REP(i, n - 1){
    int a, b; cin >> a >> b; a --; b --;
    E[a].pb(b); E[b].pb(a);
  }

  if(k == 0){
    cout << 0;
    return 0;
  }

  REP(i, n){
    for(auto e : E[i])
      sum[i] += val[e];
    sol = max(sol, sum[i]);
  }

  dfs(0);

  cout << sol;
}

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