# CEOI '20 - Star Trek

Jun 3rd, 2023
1,232
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <iostream>
2. #include <vector>
3. #include <algorithm>
4. // #include <cassert>
5.
6. using namespace std;
7.
8. const int BASE = 1e9 + 7;
9.
10. int add(int x, int y) {
11.     x += y;
12.     if (x >= BASE) return x-BASE;
13.     return x;
14. }
15. int sub(int x, int y) {
16.     x -= y;
17.     if (x < 0) return x+BASE;
18.     return x;
19. }
20. int mult(int x, int y) {
21.     return (long long)x*y % BASE;
22. }
23. int bin_pow(int x, int y) {
24.     int res = 1;
25.     while (y) {
26.         if (y & 1) res = mult(res, x);
27.         y >>= 1;
28.         x = mult(x, x);
29.     }
30.     return res;
31. }
32. int inverse(int x) {
33.     return bin_pow(x, BASE-2);
34. }
35.
36. vector<vector<int>> g;
37. vector<int> S, A;
38. vector<int> S_subtree, A_subtree; // state in subtree, cnt of change nodes in subtree
39.
40. void dfsSubtree(int node, int parent = -1) {
41.     int s[2]{}, a[2]{};
42.     for (int child : g[node]) {
43.         if (child != parent) {
44.             dfsSubtree(child, node);
45.             s[S_subtree[child]]++, a[S_subtree[child]] += A_subtree[child];
46.         }
47.     }
48.     if (s[0] == 0) { // no losing child, so this node is losing
49.         S_subtree[node] = 0;
50.         A_subtree[node] = a[1]; // any losing node would make it winning
51.         ++A_subtree[node]; // this node itself could become winning
52.     } else {
53.         S_subtree[node] = 1;
54.         A_subtree[node] = (s[0] > 1 ? 0 : a[0]);
55.     }
56. }
57. void dfsOutside(int node, int parent = -1, int S_parent = 1, int A_parent = 0) {
58.     vector<pair<int, int>> d[2]{}; // directions: (node, A_direction[node])
59.     int a[2]{};
60.     if (parent != -1) d[S_parent].emplace_back(parent, A_parent), a[S_parent] += A_parent;
61.     for (int child : g[node]) {
62.         if (child != parent) {
63.             d[S_subtree[child]].emplace_back(child, A_subtree[child]);
64.             a[S_subtree[child]] += A_subtree[child];
65.         }
66.     }
67.
68.     if (d[0].size() == 0) { // no losing children, so losing node
69.         S[node] = 0;
70.         A[node] = a[1];
71.         ++A[node]; // this node itself could become winning
72.         for (int child : g[node]) {
73.             if (child != parent) {
74.                 dfsOutside(child, node, 0, A[node] - A_subtree[child]);
75.             }
76.         }
77.     } else {
78.         S[node] = 1;
79.         if (d[0].size() == 1) {
80.             A[node] = a[0];
81.             int x = d[0][0].first;
82.             for (int child : g[node]) {
83.                 if (child != parent) {
84.                     if (child == x) dfsOutside(child, node, 0, a[1]+1); // +1 since this node itself can change
85.                     else dfsOutside(child, node, 1, a[0]);
86.                 }
87.             }
88.         } else if (d[0].size() == 2) {
89.             A[node] = 0;
90.             int x = d[0][0].first, y = d[0][1].first;
91.             int ax = (x == parent ? A_parent : A_subtree[x]);
92.             int ay = (y == parent ? A_parent : A_subtree[y]);
93.             for (int child : g[node]) {
94.                 if (child != parent) {
95.                     if (child == x) dfsOutside(child, node, 1, ay);
96.                     else if (child == y) dfsOutside(child, node, 1, ax);
97.                     else dfsOutside(child, node, 1, 0);
98.                 }
99.             }
100.         } else {
101.             A[node] = 0;
102.             for (int child : g[node]) {
103.                 if (child != parent) dfsOutside(child, node, 1, 0);
104.             }
105.         }
106.     }
107. }
108.
109. struct Matrix {
110.     int N, M;
111.     vector<vector<int>> v;
112.     Matrix(){}
113.     Matrix(int N, int M): N(N), M(M) {v.assign(N, vector<int>(M));}
114.     Matrix operator * (const Matrix &m2) {
115.         // assert(M == m2.N);
116.         int K = m2.M;
117.         Matrix m(N, K);
118.         for (int i = 0; i < N; ++i) {
119.             for (int k = 0; k < K; ++k) {
120.                 for (int j = 0; j < M; ++j) {
121.                     m.v[i][k] = add(m.v[i][k], mult(v[i][j], m2.v[j][k]));
122.                 }
123.             }
124.         }
125.         return m;
126.     }
127.     Matrix operator ^ (long long int n) {
128.         // assert(N == M);
129.         Matrix m(N, N);
130.         for (int i = 0; i < N; ++i) m.v[i][i] = 1; // identity
131.         Matrix p2 = *this;
132.         while (n) {
133.             if (n & 1) m = m * p2;
134.             n >>= 1;
135.             p2 = p2 * p2;
136.         }
137.         return m;
138.     }
139.     void print() const {
140.         for (int i = 0; i < N; ++i) {
141.             for (int j = 0; j < M; ++j) cout << v[i][j] << ' ';
142.             cout << endl;
143.         }
144.     }
145. };
146.
147. int main() {
148.     ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
149.     int N; long long D; cin >> N >> D;
150.     g.resize(1+N);
151.     for (int i = 1; i <= N-1; ++i) {
152.         int u, v; cin >> u >> v;
153.         g[u].push_back(v);
154.         g[v].push_back(u);
155.     }
156.
157.     /*
158.     Let's define a losing combination over some suffix of dimensions as the number of
159.     combinations of start node and portals in that suffix that lead to a loss for the
160.     player whose turn it is at the start node.
161.     Then the number of losing combinations over a larger suffix can be calculated.
162.     */
163.
164.     /* Find A[r] for all "roots" - the number of losing nodes s.t. if that node was turned into
165.     a winning node, that would change the state of the root r */
166.     S_subtree.resize(1+N), A_subtree.resize(1+N);
167.     dfsSubtree(1);
168.     S.resize(1+N), A.resize(1+N);
169.     dfsOutside(1);
170.
171.     int cnt[2]{}, sumA[2]{};
172.     for (int i = 1; i <= N; ++i) ++cnt[S[i]], sumA[S[i]] = add(sumA[S[i]], A[i]);
173.
174.     // This is the O(D) version:
175.     // int L = 0, N2d = 1; // N^(2d)
176.     // for (int d = 0; d <= D-1; ++d) {
177.     //     L = add(mult(N2d, cnt[0]), mult(L, sub(sumA[1], sumA[0])));
178.     //     N2d = mult(N2d, mult(N, N));
179.     // }
180.
181.     // O(log D) using Matrix Exponentiation
182.     int T = sub(sumA[1], sumA[0]);
183.     Matrix m(2, 2);
184.     m.v = {{T, cnt[0]}, {0, mult(N, N)}};
185.     Matrix v(2, 1);
186.     v.v = {{0}, {1}};
187.
188.     // m^D * v
189.     Matrix m2 = m ^ D;
190.     Matrix m3 = m2 * v;
191.     int L = m3.v[0][0], N2d = m3.v[1][0];
192.
193.     int win = 0;
194.     if (S[1] == 0) win = mult(A[1], L);
195.     else win = sub(N2d, mult(A[1], L));
196.     cout << win << endl;
197.     return 0;
198. }