Submission #1301426

#TimeUsernameProblemLanguageResultExecution timeMemory
1301426Davdav1232Wells (CEOI21_wells)C++20
0 / 100
4 ms6212 KiB
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef vector<vi> vvi; 
typedef vector<vi> vvi;
const ll MOD=1e9+7;
ll fast_pow(ll a, int b){
    ll ans=1;
    while(b){
        if(b&1) ans*=a;
        ans%=MOD;
        a*=a;
        a%=MOD;
        b>>=1;
    }
    return ans;
}

vvi G;
const int MAX_N=1500000;
// int t=0;
// int tin[MAX_N], tout[MAX_N], dep[MAX_N];
 vi anc(MAX_N);

void df(int u, int p){
    // tin[u]=t++;
    anc[u]=p;
    // for(int i=1; i<14; i++){
    //     anc[u][i]=anc[anc[u][i-1]][i-1];
    // }
    for(int v : G[u]){
        if(v==p) continue;
        // dep[v]=dep[u]+1;
        df(v, u);
    }
    // tout[u]=t++;
}

// bool is_anc(int u, int v) {return (tin[u]<=tin[v] && tout[u]>=tout[v]);}

// int lca(int u, int v){
//     if(is_anc(u, v)) return u;
//     for(int r=13; r>=0; r--){
//         if(!is_anc(anc[v][r], u)) v=anc[v][r];
//     }
//     return anc[v][0];
// }

// int dist(int u, int v) return dep[u]+dep[v]-2*dep[lca(u, v)];

int main(){
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int n, k; cin>>n>>k;
    G.resize(n);
    for(int i=1; i<n; i++){
        int a, b; cin>>a>>b; G[a-1].push_back(b-1); G[b-1].push_back(a-1);
    }
    
    //find a diameter
    int dis1[n];
    int dis2[n];
    bool vis[n];
    for(int i=0; i<n; i++) vis[i]=0;
    stack<pii> dfs;
    dfs.push({0, 0});
    while(dfs.size()){
        pii u=dfs.top();
        dfs.pop();
        if(vis[u.first]) continue;
        vis[u.first]=1;
        dis1[u.first]=u.second;
        for(int v : G[u.first]) dfs.push({v, u.second+1}); 
    }
    int m_dis=-1;
    int s=-1;
    for(int i=0; i<n; i++){
        if(dis1[i]>m_dis){
            s=i;
            m_dis=dis1[i];
        }
    }
    for(int i=0; i<n; i++) vis[i]=0;
    dfs.push({s, 0});
    while(dfs.size()){
        pii u=dfs.top();
        dfs.pop();
        if(vis[u.first]) continue;
        vis[u.first]=1;
        dis2[u.first]=u.second;
        for(int v : G[u.first]) dfs.push({v, u.second+1}); 
    }
    m_dis=-1;
    int e=-1;
    for(int i=0; i<n; i++){
        if(dis2[i]>m_dis){
            e=i;
            m_dis=dis2[i];
        }
    }
    df(e, e);
    for(int i=0; i<n; i++) vis[i]=0;
    dfs.push({e, 0});
    while(dfs.size()){
        pii u=dfs.top();
        dfs.pop();
        if(vis[u.first]) continue;
        vis[u.first]=1;
        dis1[u.first]=u.second;
        for(int v : G[u.first]) dfs.push({v, u.second+1}); 
    }
    int d=dis2[e];
    vvi H(n);
    if(d<k-1){
        cout<<"YES\n"<<fast_pow(2, n);
        return 0;
    }
    vector<int> max_dep(d+1, 0);
    for(int u=0; u<n; u++){
        int dep=dis1[u]+dis2[u]-d;
        dep/=2;
        max_dep[dis1[u]-dep]=max(max_dep[dis1[u]-dep], dep);
    }
    int pre=s;
    int sss=anc[s];
    while(sss!=e){
        if(max_dep[dis1[sss]]>=k-2){
            H[pre].push_back(anc[sss]);
            H[anc[sss]].push_back(pre);
        }
        pre=sss;
        sss=anc[sss];
    }
    vvi distances(k);
    vvi distances2(k);
    for(int i=0; i<n; i++){
        if(dis1[i]>=k || (d>=k+dis1[i])) distances[dis1[i]%k].push_back(i);
        if(dis2[i]>=k || (d>=k+dis2[i])) distances2[dis2[i]%k].push_back(i);
    }
    for(int i=0; i<k; i++){
        for(int j=1; j<distances[i].size(); j++){
            H[distances[i][j-1]].push_back(distances[i][j]);
            H[distances[i][j]].push_back(distances[i][j-1]);
        }
    }
    for(int i=0; i<k; i++){
        for(int j=1; j<distances2[i].size(); j++){
            H[distances2[i][j-1]].push_back(distances2[i][j]);
            H[distances2[i][j]].push_back(distances2[i][j-1]);
        }
    }
    vector<int> con(n, 0);
    int count=0;
    for(int i=0; i<n; i++){
        if(con[i]) continue;
        count++;
        stack<int> ss;
        ss.push(i);
        while(ss.size()){
            int u=ss.top();
            ss.pop();
            if(con[u]) continue;
            con[u]=count;
            for(int v : H[u]) ss.push(v);
        }
    }
    vector<int> clique(k);
    for(int i=0; i<n; i++){
        if(dis1[i]+dis2[i]>d) continue;
        clique[dis1[i]%k]=con[i];
    }
    vector<int> c(count+1, 0);
    for(int i=0; i<k; i++){
        c[clique[i]]++;
    }
    ll num=0;
    ll ans=1;
    for(int i=1; i<=count; i++){
        if(c[i]==0){
            ans*=2;
            ans%=MOD;
        }
        if(c[i]==1){
            num++;
        }
    }
    ans*=num;
    ans%=MOD;
    if(ans==0){
        cout<<"NO\n0";
    }
    else cout<<"YES\n"<<ans;


    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...