제출 #577356

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
577356		piOOE	Star Trek (CEOI20_startrek)	C++17	65 / 100	1067 ms	19856 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 sub(int a, int b) {
    return a >= b ? a - b : a - b + mod;
}

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);

        auto sub = canWin;

        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);
        if (sub[0]) {
            for (int i = 0; i < n; ++i) {
                if (!sub[i] && !whatIf[i] || sub[i]) {
                    ans = add(ans, n - cntWin);
                }
            }
        } else {
            for (int i = 0; i < n; ++i) {
                if (!sub[i] && whatIf[i]) {
                    ans = add(ans, n - cntWin);
                }
            }
        }
        cout << ans;
    } else {
        vector<int> num(n);
        for (int root = 0; root < n; ++root) {
            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(root, -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(root, -1);

            auto sub = canWin;

            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(root, -1);

            int cntWin = count(canWin.begin(), canWin.end(), true);
            if (sub[root]) {
                for (int i = 0; i < n; ++i) {
                    if (!sub[i] && !whatIf[i] || sub[i]) {
                        num[root] = add(num[root], 1);
                    }
                }
            } else {
                for (int i = 0; i < n; ++i) {
                    if (!sub[i] && whatIf[i]) {
                        num[root] = add(num[root], 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);
        vector<int> dp(n);
        for (int i = 0; i < n; ++i) {
            dp[i] = canWin[i];
        }
        for (int iter = 1; iter <= d; ++iter) {
            int prv = binpow(mul(n, n), iter - 1);
            int cntWin = 0;
            for (int i = 0; i < n; ++i) {
                cntWin = add(cntWin, dp[i]);
            }
            int cntLoose = sub(mul(n, prv), cntWin);
            for (int i = 0; i < n; ++i) {
                dp[i] = add(canWin[i] * mul(cntWin, n), mul(num[i], cntLoose));
            }
        }
        cout << dp[0];
    }
    return 0;
}

컴파일 시 표준 에러 (stderr) 메시지

startrek.cpp: In lambda function:
startrek.cpp:105:61: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  105 |                 } else if (!canWin[p] || cntLoosing[p] == 1 && !canWin[v]) {
startrek.cpp: In function 'int main()':
startrek.cpp:123:29: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  123 |                 if (!sub[i] && !whatIf[i] || sub[i]) {
      |                     ~~~~~~~~^~~~~~~~~~~~~
startrek.cpp: In lambda function:
startrek.cpp:191:65: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  191 |                     } else if (!canWin[p] || cntLoosing[p] == 1 && !canWin[v]) {
startrek.cpp: In function 'int main()':
startrek.cpp:208:33: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  208 |                     if (!sub[i] && !whatIf[i] || sub[i]) {
      |                         ~~~~~~~~^~~~~~~~~~~~~
startrek.cpp:205:17: warning: unused variable 'cntWin' [-Wunused-variable]
  205 |             int cntWin = count(canWin.begin(), canWin.end(), true);
      |                 ^~~~~~
startrek.cpp: In lambda function:
startrek.cpp:271:61: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  271 |                 } else if (!canWin[p] || cntLoosing[p] == 1 && !canWin[v]) {

• 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: 0 / 20
• Subtask 8: 0 / 15