Submission #660220

#	Time	Username	Problem	Language	Result	Execution time	Memory
660220		fatemetmhr	Wells (CEOI21_wells)	C++17	0 / 100	49 ms	94292 KiB

wells

// Lotfan komak


#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

#define all(x) x.begin(), x.end()
#define fi     first
#define se     second
#define pb     push_back

const int maxn5  = 2e6 + 10;
const int mod    = 1e9 + 7;

vector <int> adj[maxn5], ed[maxn5], ver;
int h[maxn5], have[maxn5], mxx = 0, mx[maxn5];
bool mark[maxn5];

void dfs(int v, int par){
    have[v]= 0;
    for(auto u : adj[v]) if(u != par){
        h[u] = h[v] + 1;
        dfs(u, v);
        have[v] = max(have[v], have[u] + 1);
    }
}

void dfs2(int v){
    mark[v] = true;
    ver.pb(v);
    for(auto u : ed[v]) if(!mark[u])
        dfs2(u);
}

void dfs_check(int v, int par){
    have[v] = 0;
    for(auto u : adj[v]) if(u != par){
        h[u] = h[v] + 1;
        dfs_check(u, v);
        if(par != -1)
            mxx = max(mxx, have[v] + have[u] + 1);
        have[v] = max(have[v], have[u] + 1);
    }

}

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

    int n, k; cin >> n >> k;
    for(int i = 0; i < n - 1; i++){
        int a, b; cin >> a >> b;
        a--; b--;
        adj[a].pb(b);
        adj[b].pb(a);
    }
    ll mul = 1;
    int dim = 0;
    for(int i = 0; i < n; i++){
        h[i] = 0;
        dfs(i, -1);
        mx[i] = 0;
        for(int j = 0; j < n; j++) if(h[j] % k == 0){
            ed[i].pb(j);
        }
        for(int j = 0; j < n; j++)
            dim = max(dim, h[j]);
        int last = 0;
        for(auto u : adj[i]){
            mx[i] = max(mx[i], have[u] + last + 1);
            last = max(last, have[u] + 1);
        }
        if(mx[i] < k - 1)
            mul = (mul * 2) % mod;
    }
    if(dim < k - 1){
        cout << "YES\n" << mul << endl;
        return 0;
    }
    ll ans = 0;
    for(int i = 0; i < n; i++) if(!mark[i] && mx[i] >= k - 1){
        ver.clear();
        dfs2(i);
        bool re = true;
        if(ver.size() == 1){
            mxx = 0;
            h[i] = 0;
            dfs_check(i, -1);
            ans += (mxx < k - 1);
            //cout << "aha " << i << ' ' << mxx << endl;
            continue;
        }
        for(auto u : ver) re &= (ed[u].size() == ver.size());
        ans += re;
    }
    ans = ans * mul % mod;
    if(ans)
        cout << "YES\n" << ans << '\n';
    else
        cout << "NO\n" << 0 << '\n';
}

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