比赛 2025.3.29 评测结果 AWWWWWWWTTTTTTTTTTTTTTTTT
题目名称 硝华流焰 最终得分 4
用户昵称 LikableP 运行时间 33.909 s
代码语言 C++ 内存使用 5.51 MiB
提交时间 2025-03-29 10:44:26
显示代码纯文本
  1. #include <cstdio>
  2.  
  3. template <typename T>
  4. T max(T x, T y) {
  5. return x > y ? x : y;
  6. }
  7.  
  8. const int MAXN = 1e5 + 10;
  9. const int MOD = 1e9 + 7;
  10.  
  11. struct EDGE {
  12. int v, w, next;
  13. } edge[MAXN << 1];
  14.  
  15. int head[MAXN], edgeNum;
  16. void AddEdge(int u, int v, int w) {
  17. edgeNum++;
  18. edge[edgeNum].v = v;
  19. edge[edgeNum].w = w;
  20. edge[edgeNum].next = head[u];
  21. head[u] = edgeNum;
  22. }
  23.  
  24. int siz[MAXN], sum;
  25. bool del[MAXN];
  26. void CalcSumAndSize(int root, int fa) {
  27. siz[root] = 1, sum++;
  28. for (int i = head[root]; i; i = edge[i].next) {
  29. int v = edge[i].v;
  30. if (v == fa || del[v]) continue;
  31. CalcSumAndSize(v, root);
  32. siz[root] += siz[v];
  33. }
  34. }
  35.  
  36. int centroid, depth[MAXN];
  37. void CalcCentroid(int root, int fa) {
  38. int maxx = sum - siz[root];
  39. depth[root] = depth[fa] + 1;
  40. for (int i = head[root]; i; i = edge[i].next) {
  41. int v = edge[i].v;
  42. if (v == fa || del[v]) continue;
  43. CalcCentroid(v, root);
  44. maxx = max(maxx, siz[v]);
  45. }
  46. if (maxx <= sum / 2) {
  47. centroid = root;
  48. }
  49. }
  50.  
  51. int zero[MAXN], one[MAXN], two[MAXN], zerodepth[MAXN], onedepth[MAXN], twodepth[MAXN], zerocnt, onecnt, twocnt;
  52. int a[MAXN], ans;
  53. void Count(int root, int fa, int c, int nzero, int none, int ntwo) {
  54. c %= MOD;
  55. if (a[root] == 0) {
  56. nzero++;
  57. } else if (a[root] == 1) {
  58. none++;
  59. } else if (a[root] == 2) {
  60. ntwo++;
  61. }
  62. if (nzero && none && ntwo) { // no missing
  63. for (int i = 1; i <= zerocnt; ++i) {
  64. (ans += zero[i] + c) %= MOD;
  65. }
  66. for (int i = 1; i <= onecnt; ++i) {
  67. (ans += one[i] + c) %= MOD;
  68. }
  69. for (int i = 1; i <= twocnt; ++i) {
  70. (ans += two[i] + c) %= MOD;
  71. }
  72. } else if ((!nzero && zerocnt) && none && ntwo) { // miss 0
  73. for (int i = 1; i <= zerocnt; ++i) {
  74. (ans += zero[i] + c) %= MOD;
  75. }
  76. for (int i = 1; i <= onecnt; ++i) {
  77. if (onedepth[i] > zerodepth[zerocnt]) {
  78. (ans += one[i] + c) %= MOD;
  79. }
  80. }
  81. for (int i = 1; i <= twocnt; ++i) {
  82. if (twodepth[i] > zerodepth[zerocnt]) {
  83. (ans += two[i] + c) %= MOD;
  84. }
  85. }
  86. } else if (nzero && (!none && onecnt) && ntwo) { // miss 1
  87. for (int i = 1; i <= zerocnt; ++i) {
  88. if (zerodepth[i] > onedepth[onecnt]) {
  89. (ans += zero[i] + c) %= MOD;
  90. }
  91. }
  92. for (int i = 1; i <= onecnt; ++i) {
  93. (ans += one[i] + c) %= MOD;
  94. }
  95. for (int i = 1; i <= twocnt; ++i) {
  96. if (twodepth[i] > onedepth[onecnt]) {
  97. (ans += two[i] + c) %= MOD;
  98. }
  99. }
  100. } else if (nzero && none && (!ntwo && twocnt)) { // miss 2
  101. for (int i = 1; i <= zerocnt; ++i) {
  102. if (zerodepth[i] > twodepth[twocnt]) {
  103. (ans += zero[i] + c) %= MOD;
  104. }
  105. }
  106. for (int i = 1; i <= onecnt; ++i) {
  107. if (onedepth[i] > twodepth[twocnt]) {
  108. (ans += one[i] + c) %= MOD;
  109. }
  110. }
  111. for (int i = 1; i <= twocnt; ++i) {
  112. (ans += two[i] + c) %= MOD;
  113. }
  114. } else if ((!nzero && zerocnt) && (!none && onecnt) && ntwo) { // miss 0 1
  115. for (int i = 1; i <= zerocnt; ++i) {
  116. if (zerodepth[i] > onedepth[1]) {
  117. (ans += zero[i] + c) %= MOD;
  118. }
  119. }
  120. for (int i = 1; i <= onecnt; ++i) {
  121. if (onedepth[i] > zerodepth[1]) {
  122. (ans += one[i] + c) %= MOD;
  123. }
  124. }
  125. for (int i = 1; i <= twocnt; ++i) {
  126. if (twodepth[i] > max(zerodepth[zerocnt], onedepth[onecnt])) {
  127. (ans += two[i] + c) %= MOD;
  128. }
  129. }
  130. } else if ((!nzero && zerocnt) && none && (!ntwo && twocnt)) { // miss 0 2
  131. for (int i = 1; i <= zerocnt; ++i) {
  132. if (zerodepth[i] > twodepth[1]) {
  133. (ans += zero[i] + c) %= MOD;
  134. }
  135. }
  136. for (int i = 1; i <= onecnt; ++i) {
  137. if (onedepth[i] > max(zerodepth[zerocnt], twodepth[zerocnt])) {
  138. (ans += one[i] + c) %= MOD;
  139. }
  140. }
  141. for (int i = 1; i <= twocnt; ++i) {
  142. if (twodepth[i] > zerodepth[1]) {
  143. (ans += two[i] + c) %= MOD;
  144. }
  145. }
  146. } else if (nzero && (!none && onecnt) && (!ntwo && twocnt)) { // miss 1 2
  147. for (int i = 1; i <= zerocnt; ++i) {
  148. if (zerodepth[i] > max(onedepth[onecnt], twodepth[twocnt])) {
  149. (ans += zero[i] + c) %= MOD;
  150. }
  151. }
  152. for (int i = 1; i <= onecnt; ++i) {
  153. if (onedepth[i] > twodepth[1]) {
  154. (ans += one[i] + c) %= MOD;
  155. }
  156. }
  157. for (int i = 1; i <= twocnt; ++i) {
  158. if (twodepth[i] > onedepth[1]) {
  159. (ans += two[i] + c) %= MOD;
  160. }
  161. }
  162. }
  163.  
  164. for (int i = head[root]; i; i = edge[i].next) {
  165. int v = edge[i].v, w = edge[i].w;
  166. if (v == fa || del[v]) continue;
  167. Count(v, root, c + w, nzero, none, ntwo);
  168. }
  169. }
  170.  
  171. void Update(int root, int fa, int c) {
  172. c %= MOD;
  173. if (a[root] == 0) {
  174. zero[++zerocnt] = c;
  175. zerodepth[zerocnt] = depth[root];
  176. } else if (a[root] == 1) {
  177. one[++onecnt] = c;
  178. onedepth[onecnt] = depth[root];
  179. } else {
  180. two[++twocnt] = c;
  181. twodepth[twocnt] = depth[root];
  182. }
  183. for (int i = head[root]; i; i = edge[i].next) {
  184. int v = edge[i].v, w = edge[i].w;
  185. if (v == fa || del[v]) continue;
  186. Update(v, root, c + w);
  187. }
  188. }
  189.  
  190. void Solve(int root) {
  191. sum = 0;
  192. CalcSumAndSize(root, 0);
  193. CalcCentroid(root, 0);
  194. root = centroid;
  195. del[root] = 1;
  196. zerocnt = onecnt = twocnt = 0;
  197. if (a[root] == 0) {
  198. zero[zerocnt = 1] = 0;
  199. zerodepth[1] = 1;
  200. } else if (a[root] == 1) {
  201. one[onecnt = 1] = 0;
  202. onedepth[1] = 1;
  203. } else {
  204. two[twocnt = 1] = 0;
  205. twodepth[1] = 1;
  206. }
  207. for (int i = head[root]; i; i = edge[i].next) {
  208. int v = edge[i].v, w = edge[i].w;
  209. if (del[v]) continue;
  210. Count(v, root, w, 0, 0, 0);
  211. Update(v, root, w);
  212. }
  213. for (int i = head[root]; i; i = edge[i].next) {
  214. int v = edge[i].v;
  215. if (del[v]) continue;
  216. Solve(v);
  217. }
  218. }
  219.  
  220. int n;
  221.  
  222. int main() {
  223. freopen("blossom.in", "r", stdin);
  224. freopen("blossom.out", "w", stdout);
  225. scanf("%d", &n);
  226. for (int i = 1; i <= n; ++i) {
  227. scanf("%d", &a[i]);
  228. }
  229. for (int i = 1; i <= n - 1; ++i) {
  230. int u, v, w;
  231. scanf("%d %d %d", &u, &v, &w);
  232. AddEdge(u, v, w);
  233. AddEdge(v, u, w);
  234. }
  235. Solve(1);
  236. printf("%d\n", ans % MOD);
  237. return 0;
  238. }