제출 #286126

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
286126		Pino	Chase (CEOI17_chase)	C++14	20 / 100	4090 ms	2304 KiB

chase

#pragma GCC optimize("Ofast,unroll-loops,no-stack-protector")
#include <bits/stdc++.h>
using namespace std;

#define df(b, e) ((b) > (e))
#define fore(i, b, e) for (auto i = (b) - df(b, e); i != e - df(b, e); i += 1 - 2 * df(b, e))
#define sz(x) (int)(x.size())
#define all(x) begin(x), end(x)
#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define mp make_pair

using lli = long long;
using ld = long double;
using ii = pair<int, int>;
using vi = vector<int>;

#ifdef LOCAL
  #include "debug.h"
#else
  #define debug(...)
#endif
const int N = 1e3 + 5;
const int M = 101;
const int INF = 1e9 + 7;

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

vi g[N];
lli dp[N][M][2], p[N], gans;
bool done[N][M][2];
int k;

template<class T>
void umax(T &a, T b) {
  a = max(a, b);
}
// bd -> breadcrumbs
// 1 <- Drop on parent's statue
// 0 <- No Drop on parent's statue

#define state [u][bc][type]
lli dfs (int u, int bc, bool type, int parent = 0) {
  auto &ans = dp state;
  if (!done state) {
    done state = 1;
    lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
      return tot + p[w]; 
    }) + p[u];
    for (auto v: g[u]) {
      if (v != parent) {
        if (bc) {
          umax(ans, res + dfs(v, bc - 1, 1, u) - p[parent] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
          umax(ans, dfs(v, bc, 0, u) - (type ? 0 : p[u]));
        }
      }
    }
    debug(u, bc, type, ans, res);
  }
  return ans;
}

lli f (int u, int croot, int bc, bool type) {
  lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
      return tot + p[w]; 
    }) + p[u];
  lli ans = 0;
  for (int v: g[u]) {
    if (v != croot) {
      if (bc) {
        umax(ans, res + dfs(v, bc - 1, 1, u) - p[croot] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
        umax(ans, dfs(v, bc, 0, u) - (type ? 0 : p[u]));
      }
    }
  }
  dp state = ans;
  done state = 1;
  return ans;
}

void solve (int u, int parent = 0) {
  lli res = accumulate(all(g[u]), 0LL, [&](lli tot, int w) -> lli {
      return tot + p[w]; 
    }) + p[u];
  lli ans = 0;
  for (int v: g[u]) {
    if (v != parent) {
     // umax(ans, res + dfs(v, k - 1, 1, u) - p[parent] - (!type ? 0 : p[u]) - (type ? 0 : p[u]));
    //  dp[u][k][0] = max(res + dfs(v, k - 1, 1, u) - p[u];
      umax(ans, res + dfs(v, k - 1, 1, u) - p[u]);
      umax(ans, dfs(v, k, 0, u) - p[u]); 
    } else {
      umax(ans, res + f(parent, u, k - 1, 1) - p[u]);
      umax(ans, f(parent, u, k, 0) - p[u]);
    }
  }
  debug(u, ans, dp[u][k][1]);
  dp[u][k][0] = ans;
  umax(gans, ans);

  for (int v: g[u]) {
    if (v != parent) {
      solve(v, u);
    }
  }
}

int main() {
  cin.tie(0) -> sync_with_stdio(0), cout.tie(0);
  #ifdef LOCAL
	  auto start_time = clock();
  #endif
  int n, u, v;
  cin >> n >> k;
  fore (i, 1, n + 1)
    cin >> p[i];
  
  fore (i, 1, n) {
    cin >> u >> v;
    g[u].eb(v);
    g[v].eb(u);
  }
  if (k == 0) {
    cout << 0 << '\n';
    return 0;
  }
  //cout << dfs(6, k, 0, 0) << '\n';
  fore (u, 1, n + 1) {
    memset(dp, 0, sizeof(dp));
    memset(done, 0, sizeof(done));
    umax(gans, dfs(u, k, 0, 0));
  }
  cout << gans << '\n';
  #ifdef LOCAL
	  auto end_time = clock();
	  cerr << "Execution time: " << (end_time - start_time) * (int)1e3 / CLOCKS_PER_SEC << " ms\n";
  #endif
  return 0;
}

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