제출 #1306794

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1306794		cpismylifeOwO	Star Trek (CEOI20_startrek)	C++20	38 / 100	336 ms	134916 KiB

startrek

#include <bits/stdc++.h>

using namespace std;

const long long mod = 1e9 + 7;
const int MaxN = 1e6 + 5;

int n;
long long d;
vector<int> graph[MaxN];

void Inp()
{
    cin >> n >> d;
    for (int x = 1; x < n; x++)
    {
        int u, v;
        cin >> u >> v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
}

struct Matrix
{
    vector<vector<long long>> mat;

    Matrix() = default;

    Matrix(vector<vector<long long>> _mat)
    {
        mat = _mat;
    }

    Matrix& operator = (const Matrix& other)
    {
        this->mat = other.mat;
        return *this;
    }

    Matrix operator + (const Matrix& other) const
    {
        Matrix res;
        assert((int)this->mat.size() == (int)other.mat.size());
        res.mat.resize((int)this->mat.size());
        for (int x = 0; x < (int)this->mat.size(); x++)
        {
            assert((int)this->mat[x].size() == (int)other.mat[x].size());
            res.mat[x].resize((int)other.mat[x].size());
            for (int y = 0; y < (int)other.mat[x].size(); y++)
            {
                res.mat[x][y] = (this->mat[x][y] + other.mat[x][y]) % mod;
            }
        }
        return res;
    }

    Matrix operator - (const Matrix& other) const
    {
        Matrix res;
        assert((int)this->mat.size() == (int)other.mat.size());
        res.mat.resize((int)this->mat.size());
        for (int x = 0; x < (int)this->mat.size(); x++)
        {
            assert((int)this->mat[x].size() == (int)other.mat[x].size());
            res.mat[x].resize((int)other.mat[x].size());
            for (int y = 0; y < (int)other.mat[x].size(); y++)
            {
                res.mat[x][y] = (this->mat[x][y] - other.mat[x][y]) % mod;
            }
        }
        return res;
    }

    Matrix operator * (const Matrix& other) const
    {
        Matrix res;
        assert((int)this->mat[0].size() == (int)other.mat.size());
        res.mat.resize((int)this->mat.size());
        for (int x = 0; x < (int)res.mat.size(); x++)
        {
            res.mat[x].resize((int)other.mat[x].size());
            for (int y = 0; y < (int)res.mat[x].size(); y++)
            {
                res.mat[x][y] = 0;
                for (int z = 0; z < (int)this->mat[0].size(); z++)
                {
                    res.mat[x][y] = (res.mat[x][y] + this->mat[x][z] * other.mat[z][y] % mod) % mod;
                }
            }
        }
        return res;
    }
};

Matrix F[MaxN];
bool isW[MaxN];
int cntnow[MaxN];

void DFS(int u, int p)
{
    isW[u] = false;
    int cnt = 0;
    for (int v : graph[u])
    {
        if (v != p)
        {
            DFS(v, u);
            isW[u] |= !isW[v];
            cnt += !isW[v];
        }
    }
    cntnow[u] = cnt;
    F[u] = Matrix({{isW[u], 1}, {!isW[u], 0}});
    for (int v : graph[u])
    {
        if (v != p)
        {
            if (cnt == !isW[v])
            {
                F[u] = (F[u] + (Matrix({{0, 1}, {1, 0}}) * F[v]));
            }
            else
            {
                F[u] = (F[u] + (Matrix({{1, 1}, {0, 0}}) * F[v]));
            }
        }
    }
}

bool Wroot[MaxN];
Matrix PreRoot[MaxN];

void Reroot(int u, int p, bool preW, Matrix preMat)
{
    Wroot[u] = isW[u] || !preW;
    int cnt = cntnow[u] + !preW;
    PreRoot[u] = Matrix({{Wroot[u], 1}, {!Wroot[u], 0}});
    if (cnt == !preW)
    {
        PreRoot[u] = (F[u] + (Matrix({{0, 1}, {1, 0}}) * preMat));
    }
    else
    {
        PreRoot[u] = (F[u] + (Matrix({{1, 1}, {0, 0}}) * preMat));
    }
    for (int v : graph[u])
    {
        if (v != p)
        {
            if (cnt == !isW[v])
            {
                PreRoot[u] = (PreRoot[u] + (Matrix({{0, 1}, {1, 0}}) * F[v]));
            }
            else
            {
                PreRoot[u] = (PreRoot[u] + (Matrix({{1, 1}, {0, 0}}) * F[v]));
            }
        }
    }
    if (cnt <= 1)
    {
        for (int v : graph[u])
        {
            if (v != p)
            {
                if (cnt == !isW[v])
                {
                    Reroot(v, u, cnt - !isW[v], PreRoot[u] - (Matrix({{0, 1}, {1, 0}}) * F[v]));
                }
                else
                {
                    Reroot(v, u, cnt - !isW[v], PreRoot[u] - (Matrix({{1, 1}, {0, 0}}) * F[v]));
                }
            }
        }
    }
    else if (cnt == 2)
    {
        Matrix newmat = PreRoot[u];
        for (int v : graph[u])
        {
            if (v != p)
            {
                if (!isW[v])
                {
                    newmat = (newmat - (Matrix({{1, 1}, {0, 0}}) * F[v]));
                    newmat = (newmat + (Matrix({{0, 1}, {1, 0}}) * F[v]));
                }
            }
        }
        for (int v : graph[u])
        {
            if (v != p)
            {
                if (!isW[v])
                {
                    Reroot(v, u, true, newmat - (Matrix({{0, 1}, {1, 0}}) * F[v]));
                }
                else
                {
                    Reroot(v, u, true, newmat - (Matrix({{1, 1}, {0, 0}}) * F[v]));
                }
            }
        }
    }
    else
    {
        for (int v : graph[u])
        {
            if (v != p)
            {
                Reroot(v, u, true, PreRoot[u] - (Matrix({{1, 1}, {0, 0}}) * F[v]));
            }
        }
    }
}

void Exc()
{
    DFS(1, -1);
    Reroot(1, -1, true, Matrix({{0, 0}, {0, 0}}));
    Matrix change({{0, 0}, {0, 0}});
    Matrix sam({{0}, {0}});
    for (int x = 1; x <= n; x++)
    {
        sam.mat[!Wroot[x]][0]++;
        change = (change + PreRoot[x]);
    }
    if (d > 1)
    {
        d--;
        sam = change * sam;
        while (d)
        {
            if (d & 1)
            {
                sam = change * sam;
            }
            change = change * change;
            d >>= 1;
        }
    }
    long long res = (F[1].mat[0][0] * sam.mat[0][0] % mod + F[1].mat[0][1] * sam.mat[1][0] % mod) % mod;
    cout << (res % mod + mod) % mod;
}

int main()
{
    //freopen("STARTREK.INP", "r", stdin);
    //freopen("STARTREK.OUT", "w", stdout);
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    int test = 1;
    //cin >> test;
    for (int x = 1; x <= test; x++)
    {
        Inp();
        Exc();
    }
    return 0;
}

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