Submission #577333

#	Time	Username	Problem	Language	Result	Execution time	Memory
577333		piOOE	Star Trek (CEOI20_startrek)	C++17	7 / 100	2 ms	2672 KiB

startrek

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N = 100000, mod = 1e9 + 7;

int add(int a, int b) {
    return a + b < mod ? a + b : a + b - mod;
}

int mul(int a, int b) {
    return a * (ll) b % mod;
}

int binpow(int a, ll p) {
    int ans = 1;
    for (; p > 0; p >>= 1, a = mul(a, a)) {
        if (p & 1) {
            ans = mul(ans, a);
        }
    }
    return ans;
}

int n;
ll d;

vector<int> g[N];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> d;
    if (n == 2) {
        cout << binpow(4, d);
        return 0;
    }
    for (int i = 1; i < n; ++i) {
        int a, b;
        cin >> a >> b;
        --a, --b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    if (d == 1) {
        vector<bool> canWin(n), whatIf(n);
        vector<int> cntLoosing(n);

        function<void(int, int)> dfs1 = [&](int v, int p) {
            canWin[v] = false;
            for (int to: g[v]) {
                if (to != p) {
                    dfs1(to, v);
                    if (!canWin[to]) {
                        canWin[v] = true;
                    }
                }
            }
        };

        dfs1(0, -1);

        function<void(int, int)> dfs2 = [&](int v, int p) {
            cntLoosing[v] = 0;
            for (int to: g[v]) {
                if (to != p) {
                    cntLoosing[v] += !canWin[to];
                }
            }
            if (p == -1) {
                whatIf[v] = true;
            } else {
                if (canWin[v]) {
                    if (!canWin[p]) {
                        whatIf[v] = whatIf[p];
                    }
                } else {
                    if (cntLoosing[p] == 1) {
                        whatIf[v] = whatIf[p];
                    }
                }
            }
            for (int to: g[v]) {
                if (to != p) {
                    dfs2(to, v);
                }
            }
        };

        dfs2(0, -1);

        function<void(int, int)> dfs3 = [&](int v, int p) {
            cntLoosing[v] = 0;
            for (int to: g[v]) {
                if (to != p) {
                    cntLoosing[v] += !canWin[to];
                } else if (!canWin[p] || cntLoosing[p] == 1 && !canWin[v]) {
                    cntLoosing[v] += 1;
                }
            }
            canWin[v] = cntLoosing[v];
            for (int to: g[v]) {
                if (to != p) {
                    dfs3(to, v);
                }
            }
        };

        dfs3(0, -1);

        int cntWin = count(canWin.begin(), canWin.end(), true);
        int ans = canWin[0] * mul(cntWin, n);
        ans = add(ans, mul(count(whatIf.begin(), whatIf.end(), !canWin[0]), n - cntWin));
        cout << ans;
    }
    return 0;
}

Compilation message (stderr)

startrek.cpp: In lambda function:
startrek.cpp:99:61: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
   99 |                 } else if (!canWin[p] || cntLoosing[p] == 1 && !canWin[v]) {

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