Submission #1014241

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014241		MilosMilutinovic	Chase (CEOI17_chase)	C++14	100 / 100	481 ms	344476 KiB

chase

#include <bits/stdc++.h>

using namespace std;

const int N = 100100;
const int K = 105;

long long dp[N][K][2][2];

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, k;
  cin >> n >> k;
  vector<int> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  vector<vector<int>> g(n);
  for (int i = 1; i < n; i++) {
    int u, v;
    cin >> u >> v;
    --u; --v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  vector<long long> s(n);
  for (int i = 0; i < n; i++) {
    for (int j : g[i]) {
      s[i] += a[j];
    }
  }
  long long res = 0;
  function<void(int, int)> Solve = [&](int v, int pv) {
    dp[v][1][0][1] = s[v];
    dp[v][1][1][1] = s[v];
    for (int i = 1; i <= k; i++) {
      for (int d = 0; d < 2; d++) {
        for (int t = 0; t < 2; t++) {
          dp[v][i][d][t] = max(dp[v][i][d][t], dp[v][i - 1][d][t]);
        }
      }
    }
    for (int u : g[v]) {
      if (u == pv) {
        continue;
      }
      Solve(u, v);
      for (int i = 0; i <= k; i++) {
        for (int me = 0; me < 2; me++) {
          for (int he = 0; he < 2; he++) {
            {
              long long ft = (he == 0 ? 0 : -a[v]);
              res = max(res, dp[v][i][0][me] + dp[u][k - i][1][he] + ft);
            }
            {
              long long ft = (me == 0 ? 0 : -a[u]);
              res = max(res, dp[v][i][1][me] + dp[u][k - i][0][he] + ft);
            }
          }
        }
      }
      for (int i = 0; i <= k; i++) {
        for (int me = 0; me < 2; me++) {
          for (int he = 0; he < 2; he++) {
            long long ft = (me == 0 ? 0 : s[v] - a[u]);
            dp[v][i + me][0][me] = max(dp[v][i + me][0][me], dp[u][i][0][he] + ft);
          }
        }
      }
      for (int i = 0; i <= k; i++) {
        for (int me = 0; me < 2; me++) {
          for (int he = 0; he < 2; he++) {
            long long ft = (me == 0 ? 0 : s[v]);
            if (he == 1) {
              ft -= a[v];
            }
            dp[v][i + me][1][me] = max(dp[v][i + me][1][me], dp[u][i][1][he] + ft);
          }
        }
      }
      for (int i = 1; i <= k; i++) {
        for (int d = 0; d < 2; d++) {
          for (int t = 0; t < 2; t++) {
            dp[v][i][d][t] = max(dp[v][i][d][t], dp[v][i - 1][d][t]);
          }
        }
      }
    }
  };
  Solve(0, 0);
  for (int i = 0; i < n; i++) {
    for (int j = 0; j <= k; j++) {
      for (int d = 0; d < 2; d++) {
        for (int t = 0; t < 2; t++) {
          res = max(res, dp[i][j][d][t]);
        }
      }
    }
  }
  cout << res << '\n';
  return 0;
}

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