Submission #999431

#	Time	Username	Problem	Language	Result	Execution time	Memory
999431		yoav_s	Star Trek (CEOI20_startrek)	C++17	43 / 100	1088 ms	34572 KiB

startrek

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef vector<ll> v;
typedef vector<v> vv;
typedef vector<vv> vvv;
typedef pair<ll,ll> p;
typedef vector<p> vp;
typedef vector<vp> vvp;
typedef vector<vvp> vvvp;
typedef pair<ll, p> tri;
typedef vector<tri> vtri;
typedef vector<vtri> vvtri;
typedef vector<vvtri> vvvtri;
typedef vector<bool> vb;
typedef vector<vb> vvb;
typedef vector<vvb> vvvb;

#define f first
#define s second
#define pb push_back
#define eb emplace_back
#define all(v) (v).begin(),(v).end()

const ll INF = 1e18;
const ll mod = 1e9 + 7;

ll modPow(ll a, ll b)
{
    if (b == 0) return 1;
    ll res = modPow(a, b / 2);
    res = (res * res) % mod;
    if (b % 2 == 1) res = (res * a) % mod;
    return res;
}

void getWinningStates(ll i, vv &graph, vb &winning, vb &visited, v &losingChildrenCount)
{
    if (visited[i]) return;
    visited[i] = true;
    for (ll x : graph[i])
    {
        if (!visited[x])
        {
            getWinningStates(x, graph, winning, visited, losingChildrenCount);
            if (!winning[x])
            {
                winning[i] = true;
                losingChildrenCount[i]++;
            }
        }
    }
}

void winningAsRoot(ll i, vv &graph, vb &visited, v &losingChildrenCount, vb &res, ll par)
{
    if (visited[i]) return;
    visited[i] = true;
    if (par != -1)
    {
        if (losingChildrenCount[i] == 0) losingChildrenCount[par]--;
        if (losingChildrenCount[par] == 0) losingChildrenCount[i]++;
    }
    res[i] = losingChildrenCount[i] > 0;
    for (ll x : graph[i])
    {
        winningAsRoot(x, graph, visited, losingChildrenCount, res, i);
    }
    if (par != -1)
    {
        if (losingChildrenCount[par] == 0) losingChildrenCount[i]--;
        if (losingChildrenCount[i] == 0) losingChildrenCount[par]++;
    }
}

ll getSubtreeSize(ll i, vv &graph, v &res)
{
    if (res[i] != 0) return 0;
    res[i] = 1;
    for (ll x : graph[i]) res[i] += getSubtreeSize(x, graph, res);
    return res[i];
}

p operator+(p a, p b)
{
    return p(a.f + b.f, a.s + b.s);
}

void operator+=(p &a, p b)
{
    a.f += b.f; a.s += b.s;
}

p operator%(p a, ll mod)
{
    return p(a.f % mod, a.s % mod);
}

void operator%=(p &a, ll mod)
{
    a = a % mod;
}

void solve(ll i, vv &graph, vb &isWinning, vvp &res, vb &visited, v &subtreeSize, vvp &childrenS)
{
    ll N = graph.size();
    visited[i] = true;
    for (ll x : graph[i])
        if (!visited[x])
        {
            solve(x, graph, isWinning, res, visited, subtreeSize, childrenS);
            childrenS[i][0] += res[x][0];
            childrenS[i][1] += res[x][1];
        }

    ll losingChildren = 0;
    ll losingChild;
    for (ll x : graph[i])
    {
        if (!visited[x] && !isWinning[x])
        {
            losingChildren++;
            losingChild = x;
        }
    }
    if (losingChildren == 0)
    {
        res[i][1] = p(N, mod - 1) + childrenS[i][0];
        res[i][0] = p(0, 1) + childrenS[i][1];
        res[i][0] %= mod;
    }
    else if (losingChildren == 1)
    {
        res[i][1] = p((N * (subtreeSize[i] - subtreeSize[losingChild])) % mod, 0);
        res[i][1] += res[losingChild][0];
        res[i][0] = res[losingChild][1];
    }
    else
    {
        res[i][0] = p(0, 0);
        res[i][1] = p((N * subtreeSize[i]) % mod, 0);
    }
    res[i][1] %= mod;
    visited[i] = false;
}

int main()
{
    ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
    long long N, D;
    cin >> N >> D;
    vv graph(N);
    for (ll i=  0; i < N - 1; i++)
    {
        long long a, b;
        cin >> a >> b;
        a--; b--;
        graph[a].pb(b);
        graph[b].pb(a);
    }
    vp amountOfWaysToWin(N);
    for (ll i=  0; i < N; i++)
    {
        vb isWinning(N), visited(N, false);
        v losingChildrenCount(N, 0);
        getWinningStates(i, graph, isWinning, visited, losingChildrenCount);
        vb winAsRoot(N, false);
        visited = vb(N, false);
        winningAsRoot(i, graph, visited, losingChildrenCount, winAsRoot, -1);
        ll winningCount = 0, losingCount = 0;
        for (auto x : winAsRoot)
        {
            if (x) winningCount++;
            else losingCount++;
        }
        v subtreeSize(N, 0);
        getSubtreeSize(i, graph, subtreeSize);
        vvp dp(N, vp(2, p(0, 0)));
        visited = vb(N, false);
        vvp childrenS(N, vp(2, p(0, 0)));
        solve(i, graph, isWinning, dp, visited, subtreeSize, childrenS);
        amountOfWaysToWin[i] = dp[i][1];
    }
    vb isWinning(N), visited(N, false);
    v losingChildrenCount(N, 0);
    getWinningStates(0, graph, isWinning, visited, losingChildrenCount);
    vb winAsRoot(N, false); visited = vb(N, false);
    winningAsRoot(0, graph, visited, losingChildrenCount, winAsRoot, -1);
    ll winningCount = 0, losingCount = 0;
    for (auto x : winAsRoot)
    {
        if (x) winningCount++;
    }
    p sum = p(0, 0);
    for (auto x : amountOfWaysToWin) sum += x;
    sum %= mod;
    ll weighted = winningCount, absolute = 1;
    for (ll i = 1; i < D; i++)
    {
        weighted = sum.f * absolute + sum.s * weighted;
        weighted %= mod;
        absolute *= (N * N) % mod;
        absolute %= mod;
    }
    cout << (weighted * amountOfWaysToWin[0].s + absolute * amountOfWaysToWin[0].f) % mod << "\n";
    return 0;
}

Compilation message (stderr)

startrek.cpp: In function 'int main()':
startrek.cpp:191:26: warning: unused variable 'losingCount' [-Wunused-variable]
  191 |     ll winningCount = 0, losingCount = 0;
      |                          ^~~~~~~~~~~

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