Submission #1156525

#	Time	Username	Problem	Language	Result	Execution time	Memory
1156525		mmdrzada	Star Trek (CEOI20_startrek)	C++20	100 / 100	55 ms	34628 KiB

startrek

#include <bits/stdc++.h>

#define int long long
using namespace std;


#define fastIO ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
#define pb push_back
#define pii pair<int, int>
#define S second
#define F first
#define sep ' '
#define vi vector<int>
#define mid (l+r)/2
#define pbb pair<bool, bool>
#define trint array<int, 3>
const int N = 1e5+10, M = 1e9+7;
int n, d;
vi adj[N];
int cnt[2][2];
int dp[N][3][2];
int sz[N];

int dp_up[N][3][2], dp_dw[N][3][2];

struct Mat {
    int n, m;
    int mat[5][5];
    Mat () {
        memset(mat, 0, sizeof mat);
    }
};
Mat mul(Mat a, Mat b) {
    assert(a.m == b.n);
    Mat res;
    res.n = a.n, res.m = b.m;
    for(int i = 0 ; i < res.n ; i ++) {
        for(int j = 0 ; j < res.m ; j ++) {
            for(int k = 0 ; k < a.m ; k ++) {
                res.mat[i][j] = (res.mat[i][j] + 1ll*a.mat[i][k]*b.mat[k][j]%M)%M;
            }
        }
    }
    return res;
}
Mat mat_pow(Mat x, int y) {
    assert(x.n == x.m);
    Mat res; res.n = x.n, res.m = x.m;
    for(int i = 0 ; i < res.n ; i ++) res.mat[i][i] = 1;
    while(y) {
        if (y&1) {
            res = mul(res, x);
        }
        y /= 2;
        x = mul(x, x);
    }
    return res;
}
int my_pow(int x, int y) {
    int res = 1;
    while(y) {
        if (y&1) res = 1ll*res*x%M;
        y /= 2;
        x = 1ll*x*x%M;
    }
    return res;
}
int inv(int x) {
	return my_pow(x, M-2);
}

/**
 * 0 -> nothing
 * 1 -> lose sued
 * 2 -> win used
 */

void gfs1(int v = 0, int p = -1) {
    if (p != -1) adj[v].erase(find(adj[v].begin(), adj[v].end(), p));
    dp_dw[v][0][0] = 1;
    sz[v]++;
    for(int neigh: adj[v]) {
        gfs1(neigh, v);
        sz[v] += sz[neigh];
        dp_dw[v][0][0] &= dp_dw[neigh][0][1];
    }
    dp_dw[v][0][1] = 1 - dp_dw[v][0][0];
    
    int s = 0;
    for(int i = 0 ; i < adj[v].size() ; i ++) {
        s += dp_dw[adj[v][i]][0][1];
    }
    dp_dw[v][2][0] += dp_dw[v][0][0]; // add the win ticket for your self

    for(int i = 0 ; i < adj[v].size() ; i ++) {
        int res1 = ((s - dp_dw[adj[v][i]][0][1]) == adj[v].size() - 1) * dp_dw[adj[v][i]][1][1];
        int res3 = ((s - dp_dw[adj[v][i]][0][1]) == adj[v].size() - 1) * dp_dw[adj[v][i]][2][1];
        dp_dw[v][1][0] += res1;
        dp_dw[v][2][0] += res3;
    }
    dp_dw[v][1][1] = sz[v] - dp_dw[v][1][0];
    dp_dw[v][2][1] = sz[v] - dp_dw[v][2][0];
}

void gfs2(int v = 0, int p = -1) {
    int s = dp_up[v][0][1];
    int sum1 = dp_up[v][1][1], sum2 = dp_up[v][2][1];
    vi zeros;
    for(int i = 0 ; i < adj[v].size() ; i ++) {
        if (dp_dw[adj[v][i]][0][1] == 0) zeros.pb(adj[v][i]);
        s += dp_dw[adj[v][i]][0][1];
        sum1 += dp_dw[adj[v][i]][1][1];
        sum2 += dp_dw[adj[v][i]][2][1];
    }

    if (s == adj[v].size()+1) {
        for(int neigh: adj[v]) {
            dp_up[neigh][0][0] = 1;
            dp_up[neigh][1][0] = sum1 - dp_dw[neigh][1][1];
            dp_up[neigh][2][0] = 1 + sum2 - dp_dw[neigh][2][1];
        }
    } else if (s == adj[v].size() && !zeros.empty()) {
        int w = zeros.front();
        dp_up[w][0][0] = 1;
        dp_up[w][1][0] = sum1 - dp_dw[w][1][1];
        dp_up[w][2][0] = 1 + sum2 - dp_dw[w][2][1];

        for(int neigh: adj[v]) {
            if (neigh != w) {
                dp_up[neigh][0][0] = 0;
                dp_up[neigh][1][0] = dp_dw[w][1][1];
                dp_up[neigh][2][0] = dp_dw[w][2][1];
            }
        }
    } else if (s == adj[v].size()) {
        for(int neigh: adj[v]) {
            dp_up[neigh][0][0] = 0;
            dp_up[neigh][1][0] = dp_up[v][1][1];
            dp_up[neigh][2][0] = dp_up[v][2][1];
        }
    }else if (s == adj[v].size()-1) {
        int w1 = zeros.front();
        int w2 = zeros.back();
        if (w1 == w2) {
            dp_up[w1][0][0] = 0;
            dp_up[w1][1][0] = dp_up[v][1][1];
            dp_up[w1][2][0] = dp_up[v][2][1];
        } else {
            dp_up[w1][0][0] = 0;
            dp_up[w1][1][0] = dp_dw[w2][1][1];
            dp_up[w1][2][0] = dp_dw[w2][2][1];

            dp_up[w2][0][0] = 0;
            dp_up[w2][1][0] = dp_dw[w1][1][1];
            dp_up[w2][2][0] = dp_dw[w1][2][1];
        }
    }

    for(int neigh: adj[v]) {
        dp_up[neigh][0][1] = 1-dp_up[neigh][0][0];
        dp_up[neigh][1][1] = n - sz[neigh] - dp_up[neigh][1][0];
        dp_up[neigh][2][1] = n - sz[neigh] - dp_up[neigh][2][0];
        gfs2(neigh, v);
    }
    dp[v][0][0] = 1;
    dp[v][0][1] = 0;
    for(int neigh: adj[v]) {
        dp[v][0][0] &= dp_dw[neigh][0][1];
    }
    dp[v][0][0] &= dp_up[v][0][1];

    dp[v][0][1] = 1 - dp[v][0][0];

    s = 0;
    for(int i = 0 ; i < adj[v].size() ; i ++) {
        s += dp_dw[adj[v][i]][0][1];
    }
    s += dp_up[v][0][1];
    dp[v][2][0] += dp[v][0][0]; // add the win ticket for your self

    for(int i = 0 ; i < adj[v].size() ; i ++) {
        int res1 = ((s - dp_dw[adj[v][i]][0][1]) == adj[v].size()) * dp_dw[adj[v][i]][1][1];
        int res3 = ((s - dp_dw[adj[v][i]][0][1]) == adj[v].size()) * dp_dw[adj[v][i]][2][1];
        dp[v][1][0] += res1;
        dp[v][2][0] += res3;
    }
    int res1 = ((s - dp_up[v][0][1]) == adj[v].size()) * dp_up[v][1][1];
    int res3 = ((s - dp_up[v][0][1]) == adj[v].size()) * dp_up[v][2][1];
    dp[v][1][0] += res1;
    dp[v][2][0] += res3;

    dp[v][1][1] = n - dp[v][1][0];
    dp[v][2][1] = n - dp[v][2][0];
    cnt[1][0] += dp[v][1][1];
    cnt[1][1] += dp[v][2][1];

    cnt[0][0] += dp[v][1][0];
    cnt[0][1] += dp[v][2][0];

    for(int i = 0 ; i < 2 ; i++) {
        for(int j = 0 ; j < 2 ; j ++) {
            cnt[i][j] %= M;
        }
    }
}

void solve() {
    cin >> n >> d;
    for(int i = 0 ; i < n-1 ; i ++) {
        int u, v; cin >> u >> v;u --, v --; adj[u].pb(v), adj[v].pb(u);
    }

    gfs1();

    dp_up[0][0][1] = 1;

    gfs2();
    Mat base; base.n = 1, base.m = 2;
    for(int root = 0 ; root < n ; root ++) {
        base.mat[0][dp[root][0][1]]++;
    }
    Mat Z; Z.n = Z.m = 2;
    for(int i = 0 ; i < 2 ; i ++) {
        for(int j = 0 ; j < 2 ; j ++) {
            Z.mat[j][i] = cnt[i][j];
        }
    }
    base = mul(base, mat_pow(Z, d-1));
    int ans = 0;
    ans += 1ll * base.mat[0][0] * (dp[0][1][1]) % M + 1ll * base.mat[0][1] * dp[0][2][1] % M;
    ans %= M;
    cout << ans << endl;
}

int32_t main() {
    fastIO;

    solve();

    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