Submission #1155880

#	Time	Username	Problem	Language	Result	Execution time	Memory
1155880		ace5	Star Trek (CEOI20_startrek)	C++20	43 / 100	1093 ms	327680 KiB

startrek

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int maxn = 100005;
const int mod = int(1e9)+7;

ll binpow(ll a,ll b)
{
    return (b == 0 ? 1 : (b%2 == 0 ? binpow((a*a)%mod,b/2) : (a*binpow(a,b-1))%mod));
}
ll del(ll a,ll b)
{
    return (a*binpow(b,mod-2))%mod;
}
ll get_ans(ll wt,ll lt,ll qwt,ll qlt,ll D,ll cw,ll w0,ll tl0,ll tw0,ll n)
{
    ll dp[2][D+1];
    for(int i = D;i > 0;--i)
    {
        if(i == D)
        {
            dp[0][i] = cw;
            dp[1][i] = n-cw;
        }
        else
        {
            dp[0][i] = (wt * binpow(n,2*D-2*i-1) + qwt * dp[0][i+1] + qlt * dp[1][i+1])%mod;
            dp[1][i] = (lt * binpow(n,2*D-2*i-1) + qwt * dp[1][i+1] + qlt * dp[0][i+1])%mod;
        }
    }
    return (w0 * binpow(n,2*D-1) + (tl0 * dp[1][1]) + (tw0 * dp[0][1]))%mod;
}

vector<int> g[maxn];
int sz[maxn];
pair<bool,int> dp[maxn];
pair<bool,int> fdp[maxn];
int n;

void dfs_dp(int v,int p)
{
    sz[v] = 1;
    vector<int> lv;
    for(auto u:g[v])
    {
        if(u != p)
        {
            dfs_dp(u,v);
            sz[v] +=sz[u];
            if(!dp[u].first)
                lv.push_back(u);
        }
    }
    int cn = 0;
    for(auto u:g[v])
    {
        if(u != p)
        {
            if(lv.size() >= 2)
            {
                cn += sz[u];
            }
            else if(lv.size() == 1)
            {
                if(u == lv[0])
                    cn += dp[u].second;
                else
                    cn += sz[u];
            }
            else
                cn += dp[u].second;
        }
    }
    dp[v] = {lv.size(),cn + (!!(lv.size()))};
    return ;
}
void dfs_fdp(int v,int p,pair<bool,int> pdp)
{
    //cout << v << ' ' << pdp.first << ' ' << pdp.second << "\n";
    vector<int> lv;
    int sum_cnt = 0;
    for(auto u:g[v])
    {
        if(u != p)
        {
            sum_cnt += dp[u].second;
            if(!dp[u].first)
            {
                lv.push_back(u);
            }
        }
    }
    if(p != -1)
    {
        sum_cnt += pdp.second;
        if(!pdp.first)
        {
            lv.push_back(p);
        }
    }
    fdp[v].first = lv.size();
    if(fdp[v].first)
    {
        if(lv.size() == 1)
            fdp[v].second = (lv[0] == p ? sz[v] + pdp.second : n - sz[lv[0]] + dp[lv[0]].second);
        else
            fdp[v].second = n;
    }
    else
        fdp[v].second = sum_cnt;
    if(lv.size() == 0)
    {
        for(auto u:g[v])
        {
            if(u != p)
            {
                dfs_fdp(u,v,{0,sum_cnt - dp[u].second});
            }
        }
    }
    else if(lv.size() == 1)
    {
        if(lv[0] == p)
        {
            for(auto u:g[v])
            {
                if(u != p)
                {
                    dfs_fdp(u,v,{1,sz[v] - sz[u] + pdp.second});
                }
            }
        }
        else
        {
            for(auto u:g[v])
            {
                if(u != p)
                {
                    if(u != lv[0])
                        dfs_fdp(u,v,{1,fdp[v].second - sz[u]});
                    else
                        dfs_fdp(u,v,{0,sum_cnt - dp[u].second});
                }
            }
        }
    }
    else if(lv.size() == 2)
    {
        for(auto u:g[v])
        {
            if(u != p)
            {
                if(u != lv[0] && u != lv[1])
                    dfs_fdp(u,v,{1,n-sz[u]});
                else if(u == lv[0])
                {
                    dfs_fdp(u,v,{1,lv[1] == p ? sz[v] - sz[u] + pdp.second : n - sz[lv[1]] - sz[u] + dp[lv[1]].second});
                }
                else
                {
                    dfs_fdp(u,v,{1,lv[0] == p ? sz[v] - sz[u] + pdp.second : n - sz[lv[0]] - sz[u] + dp[lv[0]].second});
                }
            }
        }
    }
    else
    {
        for(auto u:g[v])
        {
            if(u != p)
            {
                dfs_fdp(u,v,{1,n-sz[u]});
            }
        }
    }
}



int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    ll D;
    cin >> n >> D;
    for(int i = 0;i < n-1;++i)
    {
        int u,v;
        cin >> u >> v;
        u--;
        v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for(int i = 0;i < n;++i)
    {
        dfs_dp(i,-1);
        fdp[i] = dp[i];
    }
    //dfs_dp(0,-1);
   // dfs_fdp(0,-1,{0,0});
    ll wt = 0,lt = 0,qwt = 0,qlt = 0,cw = 0,w0 = 0,tl0 = 0,tw0 = 0;
    if(fdp[0].first)
    {
        w0 = fdp[0].second;
        tw0 = n-fdp[0].second;
    }
    else
        tl0 = n-fdp[0].second;
   // cout << dp[0].second << ' ';
    for(int j =0;j < n;++j)
    {
       // cout << fdp[j].first << ' ' << fdp[j].second << "\n";
        wt = (wt+(fdp[j].first ? fdp[j].second : 0))%mod;
        lt = (lt+(fdp[j].first ? 0 : fdp[j].second))%mod;
        if(fdp[j].first)
            qwt = (qwt+n - fdp[j].second)%mod;
        else
            qlt = (qlt+n - fdp[j].second)%mod;
        cw += (fdp[j].first);
    }
   // cout << wt << ' ' << lt << ' ' << qwt << ' ' << qlt << ' ' << cw << ' ' << w0 << ' ' << tl0 << ' ' << tw0 << "\n";
    cout << get_ans(wt,lt,qwt,qlt,D,cw,w0,tl0,tw0,n) << "\n";
    return 0;
}

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