Submission #1087975

#	Time	Username	Problem	Language	Result	Execution time	Memory
1087975		MercubytheFirst	Wells (CEOI21_wells)	C++17	0 / 100	1 ms	348 KiB

wells

#include "bits/stdc++.h"
using namespace std;
using ll = long long;
#define endl '\n'

#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 37
#endif

vector<vector<ll> > g;
ll n, k;
vector<ll> total_mark, sub_mark, dep, group;
vector<bool> useless;

void mark_dfs(ll v, ll p, ll d) {
    dep[v] = d;
    if(d % k == 0) {
        total_mark[v]++;
    }
    sub_mark[v] = !useless[v];
    for(ll to : g[v]) {
        if(to != p) {
            mark_dfs(to, v, d + 1);
            sub_mark[v] |= sub_mark[to];
        }
    }
}

bool check_dfs(ll v, ll p, ll d) {
    bool ok = true;
    ll marked_child = 0;
    for(ll to : g[v]) {
        if(to == p) {
            continue;
        }
        ok &= check_dfs(to, v, d + 1);
        if(!ok) {
            return false;
        }
        if(sub_mark[to]) {
            marked_child++;
        }
    }
    if(marked_child >= 2) {
        if(d % k == 0) {
            ok &= true;
        }
        else if(k % 2 == 1) {
            ok &= false;
        }
        else {
            ok &= (d % k == k/2);
        }
    }
    return ok;
}

bool check(ll x) {
    if(group[x] != -1) {
        return false;
    }
    dep.assign(n + 1, -1);
    sub_mark.assign(n + 1, false);
    mark_dfs(x, -1, 0);
    bool ok = check_dfs(x, -1, 0);
    if(!ok) {
        return false;
    }
    for(ll i = 1; i <= n; ++i) {
        assert(dep[i] != -1);
        if(dep[i] % k == 0) {
            group[i] = x;
        }
    }
    return true;
}

struct P
{
    static const ll inf = 1e9 + 37;
    ll len, source;
    P() { len = -inf, source = -1; }
    P(ll _len, ll _source) : len(_len), source(_source) { } 
    bool operator<(const P& other) const {
        return len < other.len;
    }
};
vector<array<P, 2> > dp1, dp2;
void longest_path_child(ll v, ll p) {
    dp1[v][0] = P(0, v);
    for(ll to : g[v]) {
        if(to == p) {
            continue;
        }
        longest_path_child(to, v);
        dp1[v] = max({dp1[v], {dp1[v][0], P(dp1[to][0].len + 1, to)}, {P(dp1[to][0].len + 1, to), dp1[v][0]}});
    }
}

void longest_path(ll v, ll p) {
    ll par_len = dp2[p][0].len;
    if(dp2[p][0].source == v) {
        par_len = dp1[p][1].len;
    }
    dp2[v] = max({dp1[v], {dp1[v][0],P(par_len + 1, p)}, {P(par_len + 1, p), dp1[v][0]}});
    for(ll to : g[v]) {
        if(to == p) {
            continue;
        }
        longest_path(to, v);
    }
}

inline void solve(){
    cin >> n >> k;
    dp1.resize(n + 1);
    dp2.resize(n + 1);
    useless.resize(n + 1);
    group.assign(n + 1, -1);
    g.assign(n + 1, vector<ll>());
    total_mark.assign(n + 1, 0);
    for(ll i = 1; i < n; ++i) {
        ll u, v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    longest_path_child(1, -1);
    longest_path(1, 0);
    for(ll i = 1; i <= n; ++i) {
        assert(dp2[i][1].len >= 0);
        if(dp2[i][0].len + dp2[i][1].len + 1 < k) {
            useless[i] = true;
        }
    }
    const ll mod = 1e9 + 7;
    ll ans = 0;
    for(ll i = 1; i <= n; ++i) {
        if(check(i)) {
            ans++;
        }
    }

    for(ll i = 1; i <= n; ++i) {
        if(useless[i]) {
            assert(group[i] == -1);
            ans = ans * 2 % mod;
        }
    }
    cout << (ans ? "YES\n" : "NO\n") << ans << endl;
}

 
signed main(){
    #ifdef LOCAL
    freopen("test.in", "r",  stdin);
    freopen("err.txt", "w", stderr);
    #endif
    ios_base::sync_with_stdio(0);
    cin.tie(NULL); 
    // signed t; cin >> t; while(t--)
        solve();
}

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