Submission #1266555

#	Time	Username	Problem	Language	Result	Execution time	Memory
1266555		MisterReaper	Star Trek (CEOI20_startrek)	C++20	100 / 100	96 ms	21064 KiB

startrek

// File C.cpp created on 02.09.2025 at 23:43:24
#include <bits/stdc++.h>

using i64 = long long;

#ifdef DEBUG 
    #include "/home/ahmetalp/Desktop/Workplace/debug.h"
#else
    #define debug(...) void(23)
#endif

template<typename T>
T power(T a, i64 b) {
    T res {1};
    while (b) {
        if (b & 1) {
            res *= a;
        }
        a *= a;
        b >>= 1;
    }
    return res;
}

constexpr int md = int(1E9) + 7;
struct MInt {
    int val;
    MInt() : val(0) {}
    template<typename T>
    MInt(T x) {
        if (-md <= x && x < md) {
            val = int(x);
        } else {
            val = int(x % md);
        }
        if (val < 0) {
            val += md;
        }
    }
    int operator()() { return val; }
    MInt& operator+= (MInt rhs) {
        if ((val += rhs.val) >= md) {
            val -= md;
        }
        return *this;
    }
    MInt& operator-= (MInt rhs) {
        if ((val -= rhs.val) < 0) {
            val += md;
        }
        return *this;
    }
    MInt& operator*= (MInt rhs) {
        val = int(1LL * val * rhs.val % md);
        return *this;
    }
    MInt inv() {
        return power(*this, md - 2);
    }
    MInt& operator/= (MInt rhs) {
        return *this *= rhs.inv();
    }
    bool operator== (MInt rhs) const {
        return val == rhs.val;
    }
    bool operator!= (MInt rhs) const {
        return val != rhs.val;
    }
};
MInt operator+ (MInt lhs, MInt rhs) {
    return lhs += rhs;
}
MInt operator- (MInt lhs, MInt rhs) {
    return lhs -= rhs;
}
MInt operator* (MInt lhs, MInt rhs) {
    return lhs *= rhs;
}
MInt operator/ (MInt lhs, MInt rhs) {
    return lhs /= rhs;
}
std::ostream& operator<< (std::ostream& os, MInt x) {
    return os << x.val;
}
std::string to_string(MInt x) {
    return to_string(x.val);
}

using Z = MInt;

Z calc(Z x, Z y, i64 d) {
    if (d == 0) {
        return 1;
    } else if (d & 1) {
        return power(x, d) + calc(x, y, d - 1) * y;
    } else {
        Z t = calc(x, y, d / 2);
        Z xp = power(x, d / 2);
        Z yp = power(y, d / 2);
        return (xp + yp) * t - xp * yp;
    }
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int N;
    i64 D;
    std::cin >> N >> D;

    std::vector<std::vector<int>> adj(N);
    for (int i = 1; i < N; ++i) {
        int A, B;
        std::cin >> A >> B;
        --A, --B;
        adj[A].emplace_back(B);
        adj[B].emplace_back(A);
    }

    int wc = 0, lc = 0;
    std::vector<int> f(N), globf(N);
    std::vector<std::array<int, 2>> cntf(N);
    auto dfs1 = [&](auto&& self, int v, int pr) -> void {
        f[v] = 0;
        for (auto u : adj[v]) {
            if (u == pr) {
                continue;
            }
            self(self, u, v);
            cntf[v][f[u]]++;
        }
        f[v] = cntf[v][0] > 0;
    };
    dfs1(dfs1, 0, 0);

    auto dfs2 = [&](auto&& self, int v, int pr) -> void {
        ((globf[v] = f[v]) ? wc : lc)++;
        for (auto u : adj[v]) {
            if (u == pr) {
                continue;
            }
            
            cntf[v][f[u]]--;
            f[v] = cntf[v][0] > 0;
            cntf[u][f[v]]++;
            f[u] = cntf[u][0] > 0;

            self(self, u, v);

            cntf[u][f[v]]--;
            f[u] = cntf[u][0] > 0;
            cntf[v][f[u]]++;
            f[v] = cntf[v][0] > 0;
        }
    };
    dfs2(dfs2, 0, 0);

    std::vector<int> cc(N);
    cntf.assign(N, {0, 0});
    std::vector<std::array<int, 2>> res(N);
    auto dfs3 = [&](auto&& self, int v, int pr) -> void {
        f[v] = 0;
        for (auto u : adj[v]) {
            if (u == pr) {
                continue;
            }
            self(self, u, v);
            res[v][f[u]] += (f[u] == 1 ? (cntf[u][0] == 1 ? res[u][0] : 0) : res[u][0] + res[u][1]);
            cntf[v][f[u]]++;
        }
        f[v] = cntf[v][0] > 0;
        res[v][0] += (f[v] == 0);
        return;
    };
    dfs3(dfs3, 0, 0);

    auto dfs4 = [&](auto&& self, int v, int pr) -> void {
        cc[v] = (f[v] == 1 ? (cntf[v][0] == 1 ? res[v][0] : 0) : res[v][0] + res[v][1]);
        for (auto u : adj[v]) {
            if (u == pr) {
                continue;
            }

            res[v][0] -= (f[v] == 0);
            cntf[v][f[u]]--;
            res[v][f[u]] -= (f[u] == 1 ? (cntf[u][0] == 1 ? res[u][0] : 0) : res[u][0] + res[u][1]);
            f[v] = cntf[v][0] > 0;
            res[v][0] += (f[v] == 0);

            res[u][0] -= (f[u] == 0);
            cntf[u][f[v]]++;
            res[u][f[v]] += (f[v] == 1 ? (cntf[v][0] == 1 ? res[v][0] : 0) : res[v][0] + res[v][1]);
            f[u] = cntf[u][0] > 0;
            res[u][0] += (f[u] == 0);
            
            self(self, u, v);

            res[u][0] -= (f[u] == 0);
            cntf[u][f[v]]--;
            res[u][f[v]] -= (f[v] == 1 ? (cntf[v][0] == 1 ? res[v][0] : 0) : res[v][0] + res[v][1]);
            f[u] = cntf[u][0] > 0;
            res[u][0] += (f[u] == 0);

            res[v][0] -= (f[v] == 0);
            cntf[v][f[u]]++;
            res[v][f[u]] += (f[u] == 1 ? (cntf[u][0] == 1 ? res[u][0] : 0) : res[u][0] + res[u][1]);
            f[v] = cntf[v][0] > 0;
            res[v][0] += (f[v] == 0);
        }
    };
    dfs4(dfs4, 0, 0);

    Z w = 0;
    for (int i = 0; i < N; ++i) {
        if (globf[i]) {
            w += cc[i];
        } else {
            w -= cc[i];
        }
    }

    debug(wc, lc, globf, cc);

    Z vr = lc * calc(Z(N) * N, w, D - 1);

    // std::vector<Z> g(D);
    // g[0] = lc;
    // for (int i = 1; i < D; ++i) {
    //     g[i] = lc * power(Z(N), 2 * i) + g[i - 1] * w;
    // }
    // debug(g, vr);

    Z ans;
    if (globf[0]) {
        ans = power(Z(N), 2 * D) - cc[0] * vr;
    } else {
        ans = cc[0] * vr;
    }

    std::cout << ans << '\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