Submission #999424

#	Time	Username	Problem	Language	Result	Execution time	Memory
999424		yoav_s	Star Trek (CEOI20_startrek)	C++17	43 / 100	1099 ms	20392 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;
}

p solve(ll i, vv &graph, vb &isWinning, vp &res, vb &visited, bool shouldWin, v &subtreeSize)
{
    ll N = graph.size();
    assert(!visited[i]);
    visited[i] = true;
    if (shouldWin)
    {
        if (isWinning[i])
        {
            ll losingChildren = 0;
            ll losingChild;
            for (ll x : graph[i])
                if (!visited[x] && !isWinning[x])
                {
                    losingChildren++;
                    losingChild = x;
                }
            if (losingChildren == 1)
            {
                res[i].f += (N * (subtreeSize[i] - subtreeSize[losingChild])) % mod;
                res[i] += solve(losingChild, graph, isWinning, res, visited, false, subtreeSize);
                res[i] %= mod;
                return res[i];
            }
            else
            {
                res[i] = p((N * subtreeSize[i]) % mod, 0);
                return res[i];
            }
        }
        else
        {
            res[i].s += mod - 1; res[i].f += N;
            for (ll x : graph[i])
            {
                if (!visited[x])
                {
                    res[i] += solve(x, graph, isWinning, res, visited, false, subtreeSize);
                }
            }
            res[i] %= mod;
            return res[i];
        }
    }
    else
    {
        //should be losing - all children should be winning
        if (!isWinning[i])
        {
            res[i].s++;
            for (ll x : graph[i])
            {
                if (!visited[x])
                    res[i] += solve(x, graph, isWinning, res, visited, true, subtreeSize);
            }
            res[i] %= mod;
            return res[i];
        }
        else
        {
            ll losingChildren = 0, losing;
            for (ll x : graph[i])
            {
                if (!visited[x] && !isWinning[x])
                {
                    losingChildren++; losing = x;
                }
            }
            if (losingChildren > 1) return p(0, 0);
            res[i] = solve(losing, graph, isWinning, res, visited, true, subtreeSize);
            return res[i];
        }
    }
}

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);
        vp dp(N, p(0, 0));
        visited = vb(N, false);
        amountOfWaysToWin[i] = solve(i, graph, isWinning, dp, visited, true, subtreeSize);
    }
    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:221:26: warning: unused variable 'losingCount' [-Wunused-variable]
  221 |     ll winningCount = 0, losingCount = 0;
      |                          ^~~~~~~~~~~
startrek.cpp: In function 'p solve(ll, vv&, vb&, vp&, vb&, bool, v&)':
startrek.cpp:175:85: warning: 'losing' may be used uninitialized in this function [-Wmaybe-uninitialized]
  175 |             res[i] = solve(losing, graph, isWinning, res, visited, true, subtreeSize);
      |                                                                                     ^

• 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