Submission #960920

#	Time	Username	Problem	Language	Result	Execution time	Memory
960920		Ghetto	Chase (CEOI17_chase)	C++17	40 / 100	4035 ms	42248 KiB

chase

#include <bits/stdc++.h>
#pragma GCC optimize("03", "unroll-loops")
#pragma GCC target("avx2", "bmi", "bmi2", "lzcnt", "popcnt")
using namespace std;
using pii = pair<int, int>;
using lint = long long;
const int MAX_N = 1e5 + 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<int> seen[MAX_N];
vector<pii> topo_order;
void dfs(int u, int prev) {
    seen[u][prev] = 1;
    for (int i = 1; i < adj[u].size(); i++) {
        if (i == prev) continue;
        int v = adj[u][i];
        if (seen[v][adj_ind[u][i]] == 1) assert(true == false);
        if (!seen[v][adj_ind[u][i]]) dfs(v, adj_ind[u][i]);
    }
    seen[u][prev] = 2;
    topo_order.push_back({u, prev});
}

vector<lint> dp[MAX_N][2];
void precomp() {
    for (int i = 1; i <= n; i++)
        seen[i].resize(adj[i].size() + 2);
    
    for (int i = 1; i <= n; i++)
        dfs(i, 0);
    reverse(topo_order.begin(), topo_order.end());

    for (int i = 1; i <= n; i++) {
        for (int c = 0; c <= 1; c++) {
            dp[i][c].resize(adj[i].size() + 2);
        }
    }
}

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

    cin.sync_with_stdio(false);
    cin.tie(0);

    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();

    assert(topo_order.size() <= 3e5);
    for (int c = 1; c <= k; c++) {
        int p = c % 2, opp_p = (c + 1) % 2;

        for (int j = topo_order.size() - 1; j >= 0; j--) {
            int u = topo_order[j].first, prev = topo_order[j].second;

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

            lint take = 0;
            for (int i = 1; i < adj[u].size(); i++)
                if (i != prev) take = max(take, dp[adj[u][i]][opp_p][adj_ind[u][i]]);
            for (int i = 1; i < adj[u].size(); i++)
                if (i != prev) take += val[adj[u][i]];
            
            dp[u][p][prev] = max(take, leave);
        }
    }

    lint ans = 0;
    for (int i = 1; i <= n; i++) ans = max(ans, dp[i][k % 2][0]);
    cout << ans << '\n';
}

Compilation message (stderr)

chase.cpp: In function 'void dfs(int, int)':
chase.cpp:23:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   23 |     for (int i = 1; i < adj[u].size(); i++) {
      |                     ~~^~~~~~~~~~~~~~~
chase.cpp: In function 'int main()':
chase.cpp:79:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   79 |             for (int i = 1; i < adj[u].size(); i++)
      |                             ~~^~~~~~~~~~~~~~~
chase.cpp:83:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   83 |             for (int i = 1; i < adj[u].size(); i++)
      |                             ~~^~~~~~~~~~~~~~~
chase.cpp:85:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   85 |             for (int i = 1; i < adj[u].size(); i++)
      |                             ~~^~~~~~~~~~~~~~~

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