Submission #960851

#TimeUsernameProblemLanguageResultExecution timeMemory
960851GhettoChase (CEOI17_chase)C++17
40 / 100
72 ms136520 KiB
#include <bits/stdc++.h>
using namespace std;
using lint = long long;
const int MAX_N = 1e4 + 5, MAX_K = 1e2 + 5;

int n, k;
lint val[MAX_N];
vector<int> adj[MAX_N], adj_ind[MAX_N]; // adj_ind[u][i] = ind of adj list u is for adj[u][i]
void init() {
    for (int i = 1; i <= n; i++) {
        adj[i].push_back(0);
        adj_ind[i].push_back(0);
    }
}

vector<lint> dp[MAX_N][MAX_K];
vector<bool> seen[MAX_N][MAX_K];
void precomp() {
    for (int i = 1; i <= n; i++) {
        for (int c = 0; c <= k; c++) {
            dp[i][c].resize(adj[i].size() + 2);
            seen[i][c].resize(adj[i].size() + 2);
        }
    }
}

lint find_dp(int u, int c, int prev) {
    if (c == 0) return 0;
    if (seen[u][c][prev]) return dp[u][c][prev];

    lint leave = 0;
    for (int i = 1; i < adj[u].size(); i++)
        if (i != prev) leave = max(leave, find_dp(adj[u][i], c, adj_ind[u][i]));

    lint take = 0;
    for (int i = 1; i < adj[u].size(); i++)
        if (i != prev) take = max(take, find_dp(adj[u][i], c - 1, adj_ind[u][i]));
     for (int i = 1; i < adj[u].size(); i++)
        if (i != prev) take += val[adj[u][i]];

    dp[u][c][prev] = max(take, leave);
    seen[u][c][prev] = true;    
    return dp[u][c][prev];
}

int main() {
    // freopen("chase.in", "r", stdin);

    cin >> n >> k;
    init();

    for (int i = 1; i <= n; i++) cin >> val[i];
    for (int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);

        int u_ind = adj[u].size() - 1, v_ind = adj[v].size() - 1;
        adj_ind[u].push_back(v_ind);
        adj_ind[v].push_back(u_ind);
    }

    precomp();

    lint ans = 0;
    for (int i = 1; i <= n; i++) ans = max(ans, find_dp(i, k, 0));
    cout << ans << '\n';
}

Compilation message (stderr)

chase.cpp: In function 'lint find_dp(int, int, int)':
chase.cpp:32:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   32 |     for (int i = 1; i < adj[u].size(); i++)
      |                     ~~^~~~~~~~~~~~~~~
chase.cpp:36:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   36 |     for (int i = 1; i < adj[u].size(); i++)
      |                     ~~^~~~~~~~~~~~~~~
chase.cpp:38:24: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   38 |      for (int i = 1; i < adj[u].size(); i++)
      |                      ~~^~~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...