Submission #1320753

#TimeUsernameProblemLanguageResultExecution timeMemory
1320753aryanFirefighting (NOI20_firefighting)C++20
22 / 100
282 ms70528 KiB
#include<bits/stdc++.h>
using namespace std;

using i64 = long long;


int main(){
        
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n,k;
    cin >> n >> k;
    assert(k <= 30);
    vector<vector<pair<int,int>>> adj(n);
    vector<vector<int>> good(n);
    for(int i = 0;i < n - 1;i++){
        int u,v,w;
        cin >> u >> v >> w;
        u --;
        v --;
        adj[u].push_back({v,w});
        adj[v].push_back({u,w});
        if(w <= k){
            good[u].push_back(v);
            good[v].push_back(u);
        }
    }
    for(int i = 0;i < n;i++){
        sort(good[i].begin(),good[i].end());
        sort(adj[i].begin(),adj[i].end());
    }
    // for(int i = 0;i < n;i++){
    //     cout << i << " : \n";
    //     for(int &e : good[i]) cout << e << " ";
    //     cout << '\n';
    // }
    vector<vector<int>> dp(n,vector<int>(3,n));
    vector<vector<int>> ndp(n,vector<int>(3));

    function<void(int,int)> dfs = [&](int u,int p){
        int cur = 0;
        int b = 0;
        for(auto pp : adj[u]){
            int v = pp.first;
            if(v == p) continue;
            dfs(v,u);
            b += dp[v][0];
        }

        // if u is bad
        int a = 1; // u exisits
        int best = n;
        int bn = -1;
        for(auto pp : adj[u]){
            int v = pp.first;
            bool f = false;
            if(cur < (int)good[u].size() && good[u][cur] == v){
                f = true;
                cur ++;
            }
            if(v == p) continue;
            // if u exsists
            if(f) a += dp[v][1];
            else a += dp[v][0];
            // if u not exsits
            if(f){
                if(best > dp[v][2] + 1 - dp[v][0]){
                best = min(best,dp[v][2] + 1 - dp[v][0]);
                bn = v;
            }
        }
    }
        dp[u][0] = min(a,b + best);
        if(a > b + best){
            assert(bn != -1);
            ndp[u][0] = bn; // make bn exists
        }else{
            ndp[u][0] = -1; // make own
        }
        // if u is good
        // case 1 : u exsits
        // case 2 : skip u
        dp[u][1] = min(a,b);
        if(a < b){
            ndp[u][1] = -1; // exsits
        }else{
            ndp[u][1] = 0; // skip
        }
        // if u exsists
        dp[u][2] = a - 1;
    };  
    dfs(0,0);
    // for(int i = 0;i < n;i++){
    //     cout << i << " : \n";
    //     for(int j = 0;j < 2;j++){
    //         cout << ndp[i][j] << " ";
    //     }
    //     cout << '\n';
    // }
    vector<int> ans;
    function<void(int,int,int)> dfs2 = [&](int u,int p,int st){
        // cout << "N " << u << " " << st << " " << ndp[u][st] << '\n';
        if(ndp[u][st] == -1 && st <= 1){
            ans.push_back(u);
        }
        if(st == 2){
            ans.push_back(u);
        }
        int cur = 0;
        int did = true;
        for(auto pp : adj[u]){
            int v = pp.first;
            bool f = false;
            if(cur < (int)good[u].size() && good[u][cur] == v){
                f = true;
                cur ++;
            }
            if(v == p) continue;
            if(st == 0){

                if(ndp[u][0] == -1){
                    // make own good
                    if(!did){
                        ans.push_back(u);
                        did = true;
                    // cout << "ADDING " << u << '\n';

                    }
                    if(f) dfs2(v,u,1);
                    else dfs2(v,u,0);
                }else{
                    if(v == ndp[u][0]){
                        assert(f);
                        dfs2(v,u,2);
                    }else{
                        dfs2(v,u,0);
                    }
                }
            }else if(st == 1){
                if(ndp[u][1] == -1){
                    if(!did){
                        ans.push_back(u);
                        did = true;
                    // cout << "ADDING " << u << '\n';

                    }
                    if(f) dfs2(v,u,1);
                    else dfs2(v,u,0);
                }else{
                    dfs2(v,u,0);
                }
            }else{
                if(!did){
                    ans.push_back(u);
                    did = true;
                    // cout << "ADDING " << u << '\n';

                }
                if(f) dfs2(v,u,1);
                else dfs2(v,u,0);
            }
        }
    };

    // cout << dp[0][0] << " " << dp[0][2] + 1 << '\n';
    if(dp[0][0] > dp[0][2] + 1){
        dfs2(0,0,2);
    }else{
        dfs2(0,0,0);
    }
    cout << (int)ans.size() << '\n';
    assert((int)ans.size() == min(dp[0][0],dp[0][2] + 1));
    for(int &e : ans) cout << e + 1 << " ";
    cout << '\n';
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...