Submission #1088027

#	Time	Username	Problem	Language	Result	Execution time	Memory
1088027		MercubytheFirst	Wells (CEOI21_wells)	C++17	0 / 100	2 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> sub_mark, dep, group;
vector<bool> useless;

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;
    }
    friend string to_string(const P& p) {
        return "{" + to_string(p.len) +  ", " + to_string(p.source) + "}";
    }
};
vector<array<P, 2> > dp1, dp2;
void longest_path_child(ll v, ll p) {
    dp1[v] = {P(0, v), P()};
    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 = dp2[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);
    }
}


void mark_dfs(ll v, ll p, ll d) {
    dep[v] = d;
    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++;
        }
    }
    // int l = (min(dp1[v][0].len, k - d%k - 1) + min(dp1[v][1].len, k - d%k - 1)) + 1;
    // debug(v, d%k, l);
    if(d % k == 0) {
        ok &= true;
    }
    else {
        int l = (min(dp1[v][0].len, k - d%k - 1) + min(dp1[v][1].len, k - d%k - 1)) + 1;
        debug(v, d%k, l);
        ok &= (l < k);
        if(dp1[v][0].len >= k - d%k and dp1[v][1].len >= k - d%k) {
            l = (min(dp1[v][0].len, k - d%k) + min(dp1[v][1].len, k - d%k)) + 1;
            debug(v, l);
            ok &= (l == k + 1);
        }
    }
    return ok;
}

bool check(ll x) {
    if(group[x] != -1) {
        return false;
    }
    cerr << x << "\n---------\n";
    dep.assign(n + 1, -1);
    sub_mark.assign(n + 1, false);
    longest_path_child(x, -1);
    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;
}


inline void solve(){
    cin >> n >> k;
    if(n > 2e4) {
        assert(false);
    }
    dp1.resize(n + 1);
    dp2.resize(n + 1);
    useless.resize(n + 1);
    group.assign(n + 1, -1);
    g.assign(n + 1, vector<ll>());
    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)) {
            debug(i, group);
            ans++;
        }
    }
    debug(group);
    for(ll i = 1; i <= n; ++i) {
        if(useless[i]) {
            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();
}

Compilation message (stderr)

wells.cpp: In function 'bool check_dfs(ll, ll, ll)':
wells.cpp:9:20: warning: statement has no effect [-Wunused-value]
    9 | #define debug(...) 37
      |                    ^~
wells.cpp:90:9: note: in expansion of macro 'debug'
   90 |         debug(v, d%k, l);
      |         ^~~~~
wells.cpp:9:20: warning: statement has no effect [-Wunused-value]
    9 | #define debug(...) 37
      |                    ^~
wells.cpp:94:13: note: in expansion of macro 'debug'
   94 |             debug(v, l);
      |             ^~~~~
wells.cpp: In function 'void solve()':
wells.cpp:9:20: warning: statement has no effect [-Wunused-value]
    9 | #define debug(...) 37
      |                    ^~
wells.cpp:152:13: note: in expansion of macro 'debug'
  152 |             debug(i, group);
      |             ^~~~~
wells.cpp:9:20: warning: statement has no effect [-Wunused-value]
    9 | #define debug(...) 37
      |                    ^~
wells.cpp:156:5: note: in expansion of macro 'debug'
  156 |     debug(group);
      |     ^~~~~

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