Submission #787946

#	Time	Username	Problem	Language	Result	Execution time	Memory
787946		definitelynotmee	Wells (CEOI21_wells)	C++17	0 / 100	19 ms	352 KiB

wells

#include<bits/stdc++.h>
#define ff first
#define ss second
#define all(x) x.begin(), x.end()
using namespace std;
using ll = long long;
using pii = pair<int,int>;
using pll = pair<ll,ll>;
template<typename T>
using matrix = vector<vector<T>>;
const int MOD = 1e9+7;
const int INF = 1e9;

int main(){
    cin.tie(0)->sync_with_stdio(0);

    int n, k;

    cin >> n >> k;

    if(n > 10000)
        return 0;

    k--;

    matrix<int> g(n+1);

    for(int i = 1; i < n; i++){
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }

    vector<ll> power(n+1);
    power[0] = 1;

    for(int i = 1; i <= n; i++)
        power[i] = power[i-1]*2%MOD;


    vector<int> d(n+1), d2(n+1);

    auto getd =[&](int id, int last, vector<int>& D, auto dfs) -> void {
        for(int i : g[id]){
            if(i != last){
                D[i] = D[id]+1;
                dfs(i,id,D,dfs);
            }
        }
    };

    getd(1,0,d,getd);

    int ponta = 1;


    for(int i = 2; i <= n; i++){
        ponta = min(ponta,i,[&](int a, int b){
            return d[a] > d[b];
        });
    }

    getd(ponta,0,d2,getd);
    
    ponta = 1;

    for(int i = 2; i <= n; i++){
        ponta = min(ponta,i,[&](int a, int b){
            return d2[a] > d2[b];
        });
    }


    fill(all(d),0);

    getd(ponta,0,d,getd);


    vector<int> dead(n+1);

    vector<int> deg(n+1);

    vector<int> st;

    for(int i = 1; i <= n; i++){
        deg[i] = g[i].size();
        if(deg[i] == 1)
            st.push_back(i);
    }

    while(!st.empty()){
        int cur = st.back();
        st.pop_back();

        if(d[cur] >= k || d2[cur] >= k)
            continue;
        
        dead[cur] = 1;

        for(int i : g[cur]){
            deg[i]--;
            if(deg[i] == 1)
                st.push_back(i);
        }
    }


    ll resp = power[accumulate(all(dead),0)];

    if(accumulate(all(dead),0) == n){
        cout << "YES\n";
        cout << resp << '\n';
        return 0;
    }

    matrix<int> g2(n+1);

    for(int i = 1; i <= n; i++){
        if(dead[i])
            continue;
        for(int j : g[i]){
            if(!dead[j])
                g2[i].push_back(j);
        }
    }

    g.swap(g2);

    k++;

    vector<int> starts;

    vector<int> startid(n+1,-1);

    for(int i = 0; i < k; i++){
        starts.push_back(ponta);

        sort(all(g[ponta]), [&](int a, int b){
            return d2[a] > d2[b];
        });
        startid[ponta] = i;

        ponta = g[ponta].back(); // menor distancia pra outra ponta

    }

    vector<int> marc(n+1);

    auto dfs =[&](int id, int last, int depth, auto dfs) -> void {
        marc[id] = depth%k == 0;
        for(int i : g[id])
            if(i != last)
                dfs(i,id,depth+1,dfs);
    };

    auto check =[&](int id, int last, int depth, int choosen, auto dfs) -> bool {
        choosen+=marc[id];
        if(depth == k-1){
            return choosen == 1;
        }
        bool ret = 1;
        for(int i : g[id]){
            if(i!=last)
                ret&=dfs(i,id,depth+1,choosen,dfs);
        }
        return ret;
    };
    int ct = 0;

    for(int start : starts){
        fill(all(marc),0);
        dfs(start,-1,0,dfs);
        bool ok = 1;

        for(int i = 1; i <= n; i++){
            ok&=check(i,-1,0,0,check);
        }
        // cerr << "->" <<  start << ' ' << ct << '\n';
        ct+=ok;

    }

    
    if(ct){
        cout << "YES\n";
        
    } else cout << "NO\n";
    
    cout << resp*ct%MOD << '\n';


}

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