Submission #999446

#	Time	Username	Problem	Language	Result	Execution time	Memory
999446		yoav_s	Star Trek (CEOI20_startrek)	C++17	100 / 100	241 ms	93008 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, vector<set<ll>> &losingChildren)
{
    if (visited[i]) return;
    visited[i] = true;
    for (ll x : graph[i])
    {
        if (!visited[x])
        {
            getWinningStates(x, graph, winning, visited, losingChildrenCount, losingChildren);
            if (!winning[x])
            {
                winning[i] = true;
                losingChildrenCount[i]++;
                losingChildren[i].insert(x);
            }
        }
    }
}

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, 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) + mod) % mod, ((a.s % mod) + mod) % 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;
}

void rerootSolve(ll i, ll par, vv &graph, vvp &res, vb &visited, v &subtreeSize, vvp &childrenS, vvp &curRootDP, vb &isWinningAsRoot, v &losingChildrenCount, vector<set<ll>> &losingChildren)
{
    ll N = graph.size();
    if (visited[i]) return;
    visited[i] = true;
    vp parChildrenS(2),myChildrenS(2),parCurRootDP(2),myCurRootDP(2);
    if (par != -1)
    {
        copy(all(childrenS[par]), parChildrenS.begin());
        copy(all(childrenS[i]), myChildrenS.begin());
        copy(all(curRootDP[par]), parCurRootDP.begin());
        copy(all(curRootDP[i]), myCurRootDP.begin());
        if (losingChildrenCount[i] == 0)
        {
            losingChildrenCount[par]--;
            losingChildren[par].erase(i);
        }
        if (losingChildrenCount[par] == 0)
        {
            losingChildrenCount[i]++;
            losingChildren[i].insert(par);
        }
        subtreeSize[par] -= subtreeSize[i];
        subtreeSize[i] += subtreeSize[par];
        childrenS[par][0] -= curRootDP[i][0];
        childrenS[par][1] -= curRootDP[i][1];
        childrenS[par][1] %= mod;
        childrenS[par][0] %= mod;
        if (losingChildrenCount[par] == 0)
        {
            curRootDP[par][1] = p(N, mod - 1) + childrenS[par][0];
            curRootDP[par][0] = p(0, 1) + childrenS[par][1];
            curRootDP[par][0] %= mod;
        }
        else if (losingChildrenCount[par] == 1)
        {
            ll losingChild = *losingChildren[par].begin();
            curRootDP[par][1] = p((N * (subtreeSize[par] - subtreeSize[losingChild])) % mod, 0);
            curRootDP[par][1] += curRootDP[losingChild][0];
            curRootDP[par][0] = curRootDP[losingChild][1];
        }
        else
        {
            curRootDP[par][0] = p(0, 0);
            curRootDP[par][1] = p((N * subtreeSize[par]) % mod, 0);
        }
        curRootDP[par][1] %= mod;
        childrenS[i][0] += curRootDP[par][0];
        childrenS[i][1] += curRootDP[par][1];
        if (losingChildrenCount[i] == 0)
        {
            curRootDP[i][1] = p(N, mod - 1) + childrenS[i][0];
            curRootDP[i][0] = p(0, 1) + childrenS[i][1];
            curRootDP[i][0] %= mod;
        }
        else if (losingChildrenCount[i] == 1)
        {
            ll losingChild = *losingChildren[i].begin();
            curRootDP[i][1] = p((N * (subtreeSize[i] - subtreeSize[losingChild])) % mod, 0);
            curRootDP[i][1] += curRootDP[losingChild][0];
            curRootDP[i][0] = curRootDP[losingChild][1];
        }
        else
        {
            curRootDP[i][0] = p(0, 0);
            curRootDP[i][1] = p((N * subtreeSize[i]) % mod, 0);
        }
        curRootDP[i][1] %= mod;
    }
    isWinningAsRoot[i] = losingChildrenCount[i] > 0;
    res[i][0] = curRootDP[i][0];
    res[i][1] = curRootDP[i][1];
    for (ll x : graph[i])
    {
        rerootSolve(x, i, graph, res, visited, subtreeSize, childrenS, curRootDP, isWinningAsRoot, losingChildrenCount, losingChildren);;
    }
    if (par != -1)
    {
        if (losingChildrenCount[par] == 0)
        {
            losingChildrenCount[i]--;
            losingChildren[i].erase(par);
        }
        if (losingChildrenCount[i] == 0)
        {
            losingChildrenCount[par]++;
            losingChildren[par].insert(i);
        }
        subtreeSize[i] -= subtreeSize[par];
        subtreeSize[par] += subtreeSize[i];
        childrenS[par] = parChildrenS;
        childrenS[i] = myChildrenS;
        curRootDP[par] = parCurRootDP;
        curRootDP[i] = myCurRootDP;
    }
}

vv multiply(vv a, vv b)
{
    vv res(a.size(), v(b[0].size()));
    for (ll i=  0; i < a.size(); i++)
    {
        for (ll j = 0; j < b[0].size(); j++)
        {
            for (ll k = 0; k < a[0].size(); k++) res[i][j] += (a[i][k] * b[k][j]) % mod;
            res[i][j] %= mod;
        }
    }
    return res;
}

vv exponent(vv a, ll exp)
{
    if (exp == 0) return {{1,0},{0,1}};
    vv res = exponent(a, exp / 2);
    res = multiply(res, res);
    if (exp % 2 == 1) res = multiply(res, a);
    return res;
}

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);
    }
    vvp amountOfWaysToGet(N, vp(2, p(0, 0)));
    vb isWinning(N), visited(N, false);
    v losingChildrenCount(N, 0);
    vector<set<ll>> losingChildren(N);
    getWinningStates(0, graph, isWinning, visited, losingChildrenCount, losingChildren);
    v subtreeSize(N, 0);
    getSubtreeSize(0, graph, subtreeSize);
    vvp dp(N, vp(2, p(0, 0)));
    visited = vb(N, false);
    vvp childrenS(N, vp(2, p(0, 0)));
    solve(0, graph, isWinning, dp, visited, subtreeSize, childrenS);
    vb winAsRoot(N, false); visited = vb(N, false);
    rerootSolve(0, -1, graph, amountOfWaysToGet, visited, subtreeSize, childrenS, dp, winAsRoot, losingChildrenCount, losingChildren);
    ll winningCount = 0, losingCount = 0;
    for (auto x : winAsRoot)
    {
        if (x) winningCount++;
    }
    p sum = p(0, 0);
    for (auto x : amountOfWaysToGet) sum += x[1];
    sum %= mod;
    ll weighted = winningCount, absolute = 1;
    vv matrix{{sum.s, sum.f},{0, (N * N) % mod}};
    vv cur = {{weighted}, {absolute}};
    vv exponented = exponent(matrix, D - 1);
    vv res = multiply(exponented, cur);
    cout << (res[0][0] * amountOfWaysToGet[0][1].s + res[1][0] * amountOfWaysToGet[0][1].f) % mod << "\n";
    return 0;
}

Compilation message (stderr)

startrek.cpp: In function 'vv multiply(vv, vv)':
startrek.cpp:240:22: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::vector<long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  240 |     for (ll i=  0; i < a.size(); i++)
      |                    ~~^~~~~~~~~~
startrek.cpp:242:26: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  242 |         for (ll j = 0; j < b[0].size(); j++)
      |                        ~~^~~~~~~~~~~~~
startrek.cpp:244:30: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  244 |             for (ll k = 0; k < a[0].size(); k++) res[i][j] += (a[i][k] * b[k][j]) % mod;
      |                            ~~^~~~~~~~~~~~~
startrek.cpp: In function 'int main()':
startrek.cpp:287:26: warning: unused variable 'losingCount' [-Wunused-variable]
  287 |     ll winningCount = 0, losingCount = 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