Submission #1213572

#	Time	Username	Problem	Language	Result	Execution time	Memory
1213572		Zero_OP	Star Trek (CEOI20_startrek)	C++20	45 / 100	58 ms	15656 KiB

startrek

#include <bits/stdc++.h>

using namespace std;

//loops (warning : long long)
#define FOR(i, l, r) for(int i = (l); i < (r); ++i)
#define ROF(i, r, l) for(int i = (r - 1); i >= l; --i)
#define rep(i, l, r) for(int i = (l); i < (r); ++i)

//pairs, tuples
#define mp make_pair
#define mt make_tuple
#define ff first
#define ss second

//vectors
#define all(v) begin(v), end(v)
#define rall(v) rbegin(v), rend(v)
#define pb push_back
#define eb emplace_back
#define sum_of(v) accumulate(all(v), 0ll)
#define sz(v) (int)v.size()
#define compact(v) v.erase(unique(all(v)), end(v))

//binary search
#define lwb lower_bound
#define upb upper_bound

//other stuffs
#define dbg(x) "[" #x " = " << (x) << "]"
#define file(task) if(fopen(task".inp", "r")){ freopen(task".inp", "r", stdin); freopen(task".out", "w", stdout); }

template<typename T>
bool minimize(T& a, const T& b){
    if(a > b) return a = b, true;
    return false;
}

template<typename T>
bool maximize(T& a, const T& b){
    if(a < b) return a = b, true;
    return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using db = double;
using pi = pair<int, int>;
using pl = pair<ll, ll>;

using vi = vector<int>;
using vb = vector<bool>;
using vl = vector<ll>;
using vpi = vector<pi>;
using vpl = vector<pl>;

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());
ll random(ll l, ll r){ return uniform_int_distribution<ll>(l,r)(rng); }

const int MAX = 1e5 + 5;

const int mod = 1e9 + 7;
template<typename T>
T ext_gcd(T a, T b, T& x, T& y){
    if(b == 0){
        x = 1, y = 0;
        return a;
    }

    T x1, y1, d = ext_gcd(b, a % b, x1, y1);
    x = y1;
    y = x1 - y1 * (a / b);
    return d;
}

struct mint{
    int v;
    mint() : v(0) {}
    mint(int v) : v(v) {
        if(v < 0) v += mod;
        if(v >= mod) v -= mod;
    }

    friend bool operator == (const mint& a, const mint& b){ return a.v == b.v; }
    friend bool operator != (const mint& a, const mint& b){ return a.v != b.v; }

    mint& operator += (const mint& o){
        v += o.v;
        if(v >= mod) v -= mod;
        return *this;
    }

    mint& operator -= (const mint& o){
        v -= o.v;
        if(v < 0) v += mod;
        return *this;
    }

    mint& operator *= (const mint& o){
        v = 1LL * v * o.v % mod;
        return *this;
    }

    mint power(ll n) const{
        mint result(1), base = *this;
        for(; n > 0; n >>= 1, base *= base){
            if(n & 1) result *= base;
        }
        return result;
    }

    mint inv() const{
        return power(mod - 2);
        //note when mod is non-prime, result = power(\varphi(mod) - 1)
    }

    mint& operator /= (const mint& o){
        return *this *= o.inv();
    }

    friend mint operator - (const mint& a){
        return mint(0) - a;
    }

    friend mint operator + (mint a, const mint& b){ return a += b; }
    friend mint operator - (mint a, const mint& b){ return a -= b; }
    friend mint operator * (mint a, const mint& b){ return a *= b; }
    friend mint operator / (mint a, const mint& b){ return a /= b; }

    friend ostream& operator << (ostream& out, const mint& v){
        return out << v.v;
    }
};

template<typename T>
struct Matrix{
    vector<vector<T>> mat;

    int row() const { return sz(mat); }
    int col() const { return sz(mat[0]); }

    auto & operator [](int i){ return mat[i]; }
    const auto & operator [](int i) const { return mat[i]; }

    Matrix() = default;
    Matrix(int n, int m, T init) : mat(n, vector<T>(m, init)) {}
    Matrix(const vector<vector<T>>& ar) : mat(ar) {}

    friend ostream& operator << (ostream& out, const Matrix& d){
        for(auto x : d.mat){
            for(auto y : x) out << y << ' ';
            out << '\n';
        }
        return out;
    }

    static Matrix identity(int n){
        Matrix a = Matrix(n, n, T(0));
        while(n--) a[n][n] = T(1);
        return a;
    }

    friend Matrix operator * (const Matrix& a, const Matrix& b){
        assert(a.col() == b.row());
        Matrix result(a.row(), b.col(), T(0));
        rep(i, 0, a.row()){
            rep(j, 0, b.col()){
                rep(k, 0, a.col()){
                    result[i][j] += a[i][k] * b[k][j];
                }
            }
        }
        return result;
    }

    Matrix& operator *= (const Matrix& b){
        return (*this) = (*this) * b;
    }

    Matrix power(long long n){
        assert(row() == col());
        Matrix base = *this, ans = identity(row());
        for(; n > 0; n >>= 1, base *= base){
            if(n & 1) ans *= base;
        }
        return ans;
    }
};

int N; ll D;
vi adj[MAX];

//Root at 1
int state[MAX], cnt[MAX];

//Reroot
int root_state[MAX], root_cnt[MAX], up_state[MAX], up_cnt[MAX];

void dfs(int u){
    cnt[u] = 1;
    for(auto v : adj[u]){
        adj[v].erase(find(all(adj[v]), u));
        dfs(v);

        if(!state[v]){
            if(!state[u]){
                cnt[u] = cnt[v];
                state[u] = 1;
            } else cnt[u] = 0;
        } else{
            if(!state[u]) cnt[u] += cnt[v];
        }
    }
}

void dfs_reroot(int u){
    int n0 = 0, n1 = 0, s0 = 0, s1 = 0;
    for(auto v : adj[u]){
        if(!state[v]) ++n0, s0 += cnt[v];
        else ++n1, s1 += cnt[v];
    }

    if(up_state[u] == 0) ++n0, s0 += up_cnt[u];
    if(up_state[u] == 1) ++n1, s1 += up_cnt[u];

    for(auto v : adj[u]){
        if(!state[v]) --n0, s0 -= cnt[v];
        else --n1, s1 -= cnt[v];

        int subtree_u_state = (n0 > 0 ? 1 : 0);
        int subtree_u_cnt = (subtree_u_state == 0 ? s1 : (n0 == 1 ? s0 : 0)) + (!subtree_u_state);

        up_state[v] = subtree_u_state;
        up_cnt[v] = subtree_u_cnt;

        // cout << dbg(v) << dbg(up_state[v]) << dbg(up_cnt[v]) << '\n';

        if(!state[v]){
            if(subtree_u_state == 0){
                root_state[v] = 1;
                root_cnt[v] = subtree_u_cnt;
            } else{
                root_state[v] = 0;
                root_cnt[v] = subtree_u_cnt + cnt[v];
                // cout << dbg(v) << dbg(cnt[v]) << dbg(subtree_u_cnt) << '\n';
            }
        } else{
            if(subtree_u_state == 0){
                root_state[v] = 1;
                root_cnt[v] = 0;
            } else{
                root_state[v] = 1;
                root_cnt[v] = cnt[v];
            }
        }

        dfs_reroot(v);
        
        if(!state[v]) ++n0, s0 += cnt[v];
        else ++n1, s1 += cnt[v];
    }
}

void testcase(int ntestcase){
    cin >> N >> D;
    if(N == 2){
        mint n2(1LL * N * N % mod);
        cout << n2.power(D) << '\n';
        return;
    }
    for(int i = 1; i < N; ++i){
        int u, v;
        cin >> u >> v;
        adj[u].pb(v);
        adj[v].pb(u);
    }

    int r = random(2, N);
    dfs(r);
    up_state[r] = -1;
    root_state[r] = state[r];
    root_cnt[r] = cnt[r];
    dfs_reroot(r);

    /*
    
    F(n) = |L| * (N^2)^d + F(n-1) * (\sum_{state_v = 1} cnt[v] - \sum_{state_v = 1} cnt[v])
    A = |L|
    B = N^2
    C = (\sum_{state_v = 1} cnt[v] - \sum_{state_v = 1} cnt[v])

    F(n) = A * (B^n) + F(n-1) * C

    [0, 1, 0 ]   [F(n-1)]  [F(n)  ]
    [0, C, AB] x [F(n) ] = [F(n+1)]
    [0, 0, B ]   [B^n]     [B^(n+1)]
    
    */

    mint A = 0, B(1LL * N * N % mod), C = 0;
    for(int i = 1; i <= N; ++i){
        A += (root_state[i] == 0);
        if(root_state[i] == 1) C += root_cnt[i];
        else C -= root_cnt[i];
    }

    Matrix<mint> trans(3, 3, 0);
    trans[0][0] = 0; trans[0][1] = 1; trans[0][2] = 0;
    trans[1][0] = 0; trans[1][1] = C; trans[1][2] = A * B;
    trans[2][0] = 0; trans[2][1] = 0; trans[2][2] = B;

    trans = trans.power(D-1);

    Matrix<mint> result(3, 1, 0);
    result[0][0] = 0;
    result[1][0] = A;
    result[2][0] = B;

    result = trans * result;
    mint sum_losing = result[1][0];

    if(!root_state[1]){
        cout << (sum_losing * root_cnt[1]) << '\n';
    } else{
        cout << (result[2][0] - sum_losing * root_cnt[1]) << '\n';
    }
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);

#ifdef LOCAL
    freopen("task.inp", "r", stdin);
    freopen("task.out", "w", stdout);
#endif // LOCAL

    int T = 1;
//    cin >> T;
    FOR(i, 0, T) testcase(i);

    return 0;
}

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