Submission #1185673

#	Time	Username	Problem	Language	Result	Execution time	Memory
1185673		Unforgettablepl	Star Trek (CEOI20_startrek)	C++20	100 / 100	521 ms	54452 KiB

startrek

#pragma GCC optimize("Ofast","unroll-loops")
#include <bits/stdc++.h>
using namespace std;

#define int long long

const int modulo = 1e9+7;


int32_t main(){
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    int N,D;
    cin >> N >> D;
    vector<vector<int>> adj(N+1);
    for(int i=1;i<N;i++){
        int a,b;cin>>a>>b;
        adj[a].emplace_back(b);
        adj[b].emplace_back(a);
    }
    vector DPbase(N+1,vector<int>(4));
    vector DPend(N+1,vector<int>(4));
    vector<int> sumzero(N+1),sumone(N+1),cntzero(N+1);
    auto calc = [&](int x){
        pair<int,int> ans;
        if(cntzero[x]){
            ans.first = 0;
            if(cntzero[x]==1)ans.second=sumzero[x];
            else ans.second = 0;
        } else {
            ans.first = 1;
            ans.second=sumone[x];
        }
        return ans;
    };
    auto modify = [&](int x,int y,int offset){
        auto t = calc(y);
        if(t.first){cntzero[x]+=offset;sumzero[x]+=t.second*offset;}
        else sumone[x]+=t.second*offset;
    };
    function<pair<int,int>(int,int)> dfs = [&](int x,int p){
        sumone[x]++;
        for(int&i:adj[x])if(i!=p){
            auto t = dfs(i,x);
            if(t.first){cntzero[x]++;sumzero[x]+=t.second;}
            else sumone[x]+=t.second;
        }
        return calc(x);
    };
    function<void(int,int)> calcDFS = [&](int x,int p){
        if(p)modify(x,p,1);
        auto t = calc(x);
        if(t.first==0){
            DPbase[x][2]=N-t.second;
            DPbase[x][0]=t.second;
            DPend[x][2]=1;
        } else {
            DPbase[x][3]=N-t.second;
            DPbase[x][1]=t.second;
            DPend[x][3]=1;
        }
        for(int&i:adj[x])if(i!=p){
            modify(x,i,-1);
            calcDFS(i,x);
            modify(x,i,1);
        }
        if(p)modify(x,p,-1);
    };
    dfs(1,0);
    calcDFS(1,0);
    DPbase[0][0]=1;
    DPend[0][0]=-1;
    auto binexpo = [&](int x,int p){
        if(p==-2)return x;
        int ans = 1;
        while(p){
            if(p&1){
                ans*=x;
                ans%=modulo;
            }
            p/=2;
            x*=x;
            x%=modulo;
        }
        return ans;
    };
    auto combine = [&](vector<vector<int>> &DP1,vector<vector<int>> &DP2){
        vector DP(N+1,vector<int>(4));
        vector<int> sums(4);
        for(int i=1;i<=N;i++){
            for(int j=0;j<4;j++){
                sums[j]+=DP2[i][j];
            }
        }
        for(int i=0;i<4;i++)sums[i]%=modulo;
        for(int i=1;i<=N;i++){
            DP[i][0]=DP1[i][1]*sums[1]+DP1[i][0]*sums[0];
            DP[i][0]%=modulo;
            DP[i][1]=DP1[i][1]*sums[0]+DP1[i][0]*sums[1];
            DP[i][1]%=modulo;
            DP[i][2]=DP1[i][2]*binexpo(N,2*(DP2[0][0]))+
                    (DP1[i][0]*sums[2]+
                     DP1[i][1]*sums[3])%modulo;
            DP[i][2]%=modulo;
            DP[i][3]=DP1[i][3]*binexpo(N,2*(DP2[0][0]))+
                    (DP1[i][1]*sums[2]+
                     DP1[i][0]*sums[3])%modulo;
            DP[i][3]%=modulo;
        }
        DP[0][0]=DP1[0][0]+DP2[0][0];
        return DP;
    };
    function<vector<vector<int>>(int)> solve = [&](int d){
        if(d==1)return DPbase;
        auto half = solve(d/2);
        auto DP = combine(half,half);
        if(d&1)DP=combine(DP,DPbase);
        return DP;
    };
    D%=modulo-1;
    auto s = solve(D);
    cout << combine(s,DPend)[1][2] << '\n';
}

• 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