Submission #286756

#	Time	Username	Problem	Language	Result	Execution time	Memory
286756		Pino	Chase (CEOI17_chase)	C++14	20 / 100	4061 ms	291864 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 = 1e5 + 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;
    ans = 0;
    if (bc) {
      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) {
            umax(ans, res + dfs(v, bc - 1, 1, u) - p[parent] - p[u]);
            umax(ans, dfs(v, bc, 0, u) - (type ? 0 : p[u]));
        }
      }
    }
  }
  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];

  vector<lli> suf, pre;
  lli ans = 0;
  for (int v: g[u]) {
    umax(ans, res + dfs(v, k - 1, 1, u) - p[u]);
    umax(ans, dfs(v, k, 0, u) - p[u]); 
    pre.eb(ans);
  }

  reverse(all(g[u]));
  ans = 0;
  for (int v: g[u]) {
    umax(ans, res + dfs(v, k - 1, 1, u) - p[u]);
    umax(ans, dfs(v, k, 0, u) - p[u]); 
    suf.eb(ans);
  }

  reverse(all(suf));
  reverse(all(g[u]));
  

  umax(gans, ans);
  debug(u, ans, u, parent);
  function<lli(int)> get = [&] (int cnt) {
    return max(cnt == 0 ? 0: pre[cnt - 1], cnt == sz(suf) - 1 ? 0: suf[cnt + 1]);
  };

  int cnt = 0;
  lli aux[M][2];
  bool ndone[M][2];
  for (int v: g[u]) {
    if (v != parent) {
      memset(done[u], 0, sizeof(done[u]));
      done[u][k][0] = 1;
      dp[u][k][0] = get(cnt);

      fore (i, 0, M)
        fore (j, 0, 2)
          aux[i][j] = dp[v][i][j];
        
      fore (i, 0, M)
        fore (j, 0, 2)
          ndone[i][j] = done[v][i][j];

      solve(v, u);

      fore (i, 0, M)
        fore (j, 0, 2)
          dp[v][i][j] = aux[i][j];
      
      fore (i, 0, M)
        fore (j, 0, 2)
          done[v][i][j] = ndone[i][j];
    }
    cnt++;
  }
  //memset(done[u], 0, sizeof(done[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;
  }
  dfs(1, k, 0, 0);
  solve(1, 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