CEOI '20 - Star Trek

#include <iostream>
#include <vector>
#include <algorithm>
// #include <cassert>

using namespace std;

const int BASE = 1e9 + 7;

int add(int x, int y) {
    x += y;
    if (x >= BASE) return x-BASE;
    return x;
}
int sub(int x, int y) {
    x -= y;
    if (x < 0) return x+BASE;
    return x;
}
int mult(int x, int y) {
    return (long long)x*y % BASE;
}
int bin_pow(int x, int y) {
    int res = 1;
    while (y) {
        if (y & 1) res = mult(res, x);
        y >>= 1;
        x = mult(x, x);
    }
    return res;
}
int inverse(int x) {
    return bin_pow(x, BASE-2);
}

vector<vector<int>> g;
vector<int> S, A;
vector<int> S_subtree, A_subtree; // state in subtree, cnt of change nodes in subtree

void dfsSubtree(int node, int parent = -1) {
    int s[2]{}, a[2]{};
    for (int child : g[node]) {
        if (child != parent) {
            dfsSubtree(child, node);
            s[S_subtree[child]]++, a[S_subtree[child]] += A_subtree[child];
        }
    }
    if (s[0] == 0) { // no losing child, so this node is losing
        S_subtree[node] = 0;
        A_subtree[node] = a[1]; // any losing node would make it winning
        ++A_subtree[node]; // this node itself could become winning
    } else {
        S_subtree[node] = 1;
        A_subtree[node] = (s[0] > 1 ? 0 : a[0]);
    }
}
void dfsOutside(int node, int parent = -1, int S_parent = 1, int A_parent = 0) {
    vector<pair<int, int>> d[2]{}; // directions: (node, A_direction[node])
    int a[2]{};
    if (parent != -1) d[S_parent].emplace_back(parent, A_parent), a[S_parent] += A_parent;
    for (int child : g[node]) {
        if (child != parent) {
            d[S_subtree[child]].emplace_back(child, A_subtree[child]);
            a[S_subtree[child]] += A_subtree[child];
        }
    }

    if (d[0].size() == 0) { // no losing children, so losing node
        S[node] = 0;
        A[node] = a[1];
        ++A[node]; // this node itself could become winning
        for (int child : g[node]) {
            if (child != parent) {
                dfsOutside(child, node, 0, A[node] - A_subtree[child]);
            }
        }
    } else {
        S[node] = 1;
        if (d[0].size() == 1) {
            A[node] = a[0];
            int x = d[0][0].first;
            for (int child : g[node]) {
                if (child != parent) {
                    if (child == x) dfsOutside(child, node, 0, a[1]+1); // +1 since this node itself can change
                    else dfsOutside(child, node, 1, a[0]);
                }
            }
        } else if (d[0].size() == 2) {
            A[node] = 0;
            int x = d[0][0].first, y = d[0][1].first;
            int ax = (x == parent ? A_parent : A_subtree[x]);
            int ay = (y == parent ? A_parent : A_subtree[y]);
            for (int child : g[node]) {
                if (child != parent) {
                    if (child == x) dfsOutside(child, node, 1, ay);
                    else if (child == y) dfsOutside(child, node, 1, ax);
                    else dfsOutside(child, node, 1, 0);
                }
            }
        } else {
            A[node] = 0;
            for (int child : g[node]) {
                if (child != parent) dfsOutside(child, node, 1, 0);
            }
        }
    }
}

struct Matrix {
    int N, M;
    vector<vector<int>> v;
    Matrix(){}
    Matrix(int N, int M): N(N), M(M) {v.assign(N, vector<int>(M));}
    Matrix operator * (const Matrix &m2) {
        // assert(M == m2.N);
        int K = m2.M;
        Matrix m(N, K);
        for (int i = 0; i < N; ++i) {
            for (int k = 0; k < K; ++k) {
                for (int j = 0; j < M; ++j) {
                    m.v[i][k] = add(m.v[i][k], mult(v[i][j], m2.v[j][k]));
                }
            }
        }
        return m;
    }
    Matrix operator ^ (long long int n) {
        // assert(N == M);
        Matrix m(N, N);
        for (int i = 0; i < N; ++i) m.v[i][i] = 1; // identity
        Matrix p2 = *this;
        while (n) {
            if (n & 1) m = m * p2;
            n >>= 1;
            p2 = p2 * p2;
        }
        return m;
    }
    void print() const {
        for (int i = 0; i < N; ++i) {
            for (int j = 0; j < M; ++j) cout << v[i][j] << ' ';
            cout << endl;
        }
    }
};

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int N; long long D; cin >> N >> D;
    g.resize(1+N);
    for (int i = 1; i <= N-1; ++i) {
        int u, v; cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    /*
    Let's define a losing combination over some suffix of dimensions as the number of
    combinations of start node and portals in that suffix that lead to a loss for the
    player whose turn it is at the start node.
    Then the number of losing combinations over a larger suffix can be calculated.
    */

    /* Find A[r] for all "roots" - the number of losing nodes s.t. if that node was turned into
    a winning node, that would change the state of the root r */
    S_subtree.resize(1+N), A_subtree.resize(1+N);
    dfsSubtree(1);
    S.resize(1+N), A.resize(1+N);
    dfsOutside(1);

    int cnt[2]{}, sumA[2]{};
    for (int i = 1; i <= N; ++i) ++cnt[S[i]], sumA[S[i]] = add(sumA[S[i]], A[i]);

    // This is the O(D) version:
    // int L = 0, N2d = 1; // N^(2d)
    // for (int d = 0; d <= D-1; ++d) {
    //     L = add(mult(N2d, cnt[0]), mult(L, sub(sumA[1], sumA[0])));
    //     N2d = mult(N2d, mult(N, N));
    // }

    // O(log D) using Matrix Exponentiation
    int T = sub(sumA[1], sumA[0]);
    Matrix m(2, 2);
    m.v = {{T, cnt[0]}, {0, mult(N, N)}};
    Matrix v(2, 1);
    v.v = {{0}, {1}};

    // m^D * v
    Matrix m2 = m ^ D;
    Matrix m3 = m2 * v;
    int L = m3.v[0][0], N2d = m3.v[1][0];

    int win = 0;
    if (S[1] == 0) win = mult(A[1], L);
    else win = sub(N2d, mult(A[1], L));
    cout << win << endl;
    return 0;
}