Submission #1002540

#	Time	Username	Problem	Language	Result	Execution time	Memory
1002540		doducanh	Star Trek (CEOI20_startrek)	C++14	100 / 100	46 ms	18008 KiB

startrek

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int mod=1e9+7;
const int maxn=1e5+7;
vector<int>a[maxn];
int n,m,A,B,L;
int lc[maxn],wc[maxn];
int dp[maxn],dpc[maxn];
int dpr[maxn],dpcrit[maxn];
int add(int a,int b){return (a+b)%mod;}
int sub(int a,int b){return (a%mod-b%mod+mod)%mod;}
int mul(int a,int b){return 1ll*a*b%mod;}
int Pow(int a,int b)
{
    int res=1;
    for(;b;b=b/2,a=(a*a)%mod)if(b&1)res=(res*a)%mod;
    return res;
}
void dfs(int u,int par)
{
    dpr[u]=0;
    for(int v:a[u])if(v!=par){
        dfs(v,u);
        if(dpr[v]==0)dpr[u]++;
    }
}
void dfs1(int u,int par)
{
    if(dpr[u]==0)dpcrit[u]=1;
    for(int v:a[u])if(v!=par){
        dfs1(v,u);
        if(!dpr[v])lc[u]+=dpcrit[v];
        else wc[u]+=dpcrit[v];
        if(dpr[u]==1&&dpr[v]==0)dpcrit[u]+=dpcrit[v];
        if(dpr[u]==0&&dpr[v]==1)dpcrit[u]+=dpcrit[v];
    }
}
void change_root(int aft,int bef)
{
    if(!dpr[aft]){
        dpr[bef]--;
        lc[bef]-=dpcrit[aft];
    }
    else wc[bef]-=dpcrit[aft];
    if(dpr[bef]>=2)dpcrit[bef]=0;
    if(dpr[bef]==1)dpcrit[bef]=lc[bef];
    if(dpr[bef]==0)dpcrit[bef]=wc[bef]+1;
    if(dpr[bef]==0){
        dpr[aft]++;
        lc[aft]+=dpcrit[bef];
    }
    else wc[aft]+=dpcrit[bef];
    if(dpr[aft]>=2){
        dpcrit[aft]=0;
    }
    if(dpr[aft]==1){
        dpcrit[aft]=lc[aft];
    }
    if(dpr[aft]==0){
        dpcrit[aft]=wc[aft]+1;
    }
}
void dfs2(int u,int par)
{
    dp[u]=dpr[u];
    dpc[u]=dpcrit[u];
    for(int v:a[u])if(v!=par){
        change_root(v,u);
        dfs2(v,u);
        change_root(u,v);
    }
}
int calc(int d)
{
    if(d==0)return L;
    if(d&1){
        return mul(calc(d/2),add(Pow(A,d/2+1),Pow(B,d/2+1)));
    }
    else{
        int X=calc(d/2-1);
        return add(mul(X,Pow(B,d/2+1)),add(mul(L,mul(Pow(A,d/2),Pow(B,d/2))),mul(X,Pow(A,d/2+1))));
    }
}
main()
{
    ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    cin>>n>>m;
    for(int i=1;i<n;i++){
        int x,y;
        cin>>x>>y;
        a[x].push_back(y);
        a[y].push_back(x);
    }
    if(n==2){
        cout<<Pow(4,m);
        return 0;
    }
    dfs(1,0);
    dfs1(1,0);
    dfs2(1,0);
//    for(int i=1;i<=n;i++){
//        cout<<i<<" "<<dp[i]<<" "<<dpc[i]<<"\n";
//    }
    L=0;
    int CW=0,CL=0;
    for(int i=1;i<=n;i++){
        if(dp[i])CW+=dpc[i];
        else L++,CL+=dpc[i];
    }
    B=sub(CW,CL);
    A=Pow(n,2);
//    cout<<add(A,B)<<"\n";
//    cout<<calc(1)<<"\n";
    int LD=calc(m-1);
//    cout<<LD<<"\n";
    if(dp[1]){
        cout<<sub(Pow(n,2*m),mul(dpc[1],LD));
    }
    else cout<<sub(Pow(n,2*m),sub(Pow(n,2*m),mul(dpc[1],LD)));

}

Compilation message (stderr)

startrek.cpp:85:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   85 | main()
      | ^~~~

• Subtask 1: 0 / 0
• Subtask 2: 7 / 7
• Subtask 3: 8 / 8
• Subtask 4: 15 / 15
• Subtask 5: 15 / 15
• Subtask 6: 20 / 20
• Subtask 7: 20 / 20
• Subtask 8: 15 / 15