Submission #1155896

#	Time	Username	Problem	Language	Result	Execution time	Memory
1155896		ace5	Star Trek (CEOI20_startrek)	C++20	100 / 100	88 ms	32908 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;
}
vector<vector<ll>> mul(vector<vector<ll>> m1,vector<vector<ll>> m2)
{
    vector<vector<ll>> ans(m1.size(),vector<ll>(m1.size()));
    for(int i = 0;i < m1.size();++i)
    {
        for(int j = 0;j < m1.size();++j)
        {
            for(int k = 0;k < m1.size();++k)
            {
                ans[i][j] += (m1[i][k] * m2[k][j])%mod;
            }
        }
    }
    return ans;
}
vector<vector<ll>> binpowm(vector<vector<ll>> a,ll b)
{
    if(b == 0)
    {
        vector<vector<ll>> ans(a.size(),vector<ll>(a.size(),0));
        for(int j =0;j < ans.size();++j)
        {
            ans[j][j] = 1;
        }
        return ans;
    }
    else
    {
        if(b%2 == 0)
        {
            vector<vector<ll>> ans = binpowm(a,b/2);
            return mul(ans,ans);
        }
        else
        {
            vector<vector<ll>> ans = binpowm(a,b-1);
            return mul(ans,a);
        }
    }
}

ll get_ans(ll wt,ll lt,ll qwt,ll qlt,ll D,ll cw,ll w0,ll tl0,ll tw0,ll n)
{
    vector<vector<ll>> mtr = {{qwt,qlt,wt,lt},{qlt,qwt,lt,wt},{0,0,(n*n)%mod,0},{0,0,0,(n*n)%mod}};
    vector<vector<ll>> ansdp = binpowm(mtr,D-1);
    ll p00 = ansdp[0][0],p01 = ansdp[0][1],p10 = ansdp[1][0],p11 = ansdp[1][1],p0w = ansdp[0][2],p1w = ansdp[1][2];
    return (w0 * binpow(n,2*D-1) + (tl0 * ((p1w * n + p10 * cw + p11 * (n-cw))%mod)) + (tw0 * ((p0w * n + p00 * cw + p01 * (n-cw))%mod)))%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);
    }
    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: 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