QOJ.ac

QOJ

IDProblemSubmitterResultTimeMemoryLanguageFile sizeSubmit timeJudge time
#699177#7997. 树 V 图hos_lyricAC ✓1593ms40964kbC++146.4kb2024-11-02 03:03:422024-11-02 03:03:43

Judging History

你现在查看的是最新测评结果

  • [2024-11-02 03:03:43]
  • 评测
  • 测评结果:AC
  • 用时:1593ms
  • 内存:40964kb
  • [2024-11-02 03:03:42]
  • 提交

answer

#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <functional>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace std;

using Int = long long;

template <class T1, class T2> ostream &operator<<(ostream &os, const pair<T1, T2> &a) { return os << "(" << a.first << ", " << a.second << ")"; };
template <class T> ostream &operator<<(ostream &os, const vector<T> &as) { const int sz = as.size(); os << "["; for (int i = 0; i < sz; ++i) { if (i >= 256) { os << ", ..."; break; } if (i > 0) { os << ", "; } os << as[i]; } return os << "]"; }
template <class T> void pv(T a, T b) { for (T i = a; i != b; ++i) cerr << *i << " "; cerr << endl; }
template <class T> bool chmin(T &t, const T &f) { if (t > f) { t = f; return true; } return false; }
template <class T> bool chmax(T &t, const T &f) { if (t < f) { t = f; return true; } return false; }
#define COLOR(s) ("\x1b[" s "m")

////////////////////////////////////////////////////////////////////////////////
template <unsigned M_> struct ModInt {
  static constexpr unsigned M = M_;
  unsigned x;
  constexpr ModInt() : x(0U) {}
  constexpr ModInt(unsigned x_) : x(x_ % M) {}
  constexpr ModInt(unsigned long long x_) : x(x_ % M) {}
  constexpr ModInt(int x_) : x(((x_ %= static_cast<int>(M)) < 0) ? (x_ + static_cast<int>(M)) : x_) {}
  constexpr ModInt(long long x_) : x(((x_ %= static_cast<long long>(M)) < 0) ? (x_ + static_cast<long long>(M)) : x_) {}
  ModInt &operator+=(const ModInt &a) { x = ((x += a.x) >= M) ? (x - M) : x; return *this; }
  ModInt &operator-=(const ModInt &a) { x = ((x -= a.x) >= M) ? (x + M) : x; return *this; }
  ModInt &operator*=(const ModInt &a) { x = (static_cast<unsigned long long>(x) * a.x) % M; return *this; }
  ModInt &operator/=(const ModInt &a) { return (*this *= a.inv()); }
  ModInt pow(long long e) const {
    if (e < 0) return inv().pow(-e);
    ModInt a = *this, b = 1U; for (; e; e >>= 1) { if (e & 1) b *= a; a *= a; } return b;
  }
  ModInt inv() const {
    unsigned a = M, b = x; int y = 0, z = 1;
    for (; b; ) { const unsigned q = a / b; const unsigned c = a - q * b; a = b; b = c; const int w = y - static_cast<int>(q) * z; y = z; z = w; }
    assert(a == 1U); return ModInt(y);
  }
  ModInt operator+() const { return *this; }
  ModInt operator-() const { ModInt a; a.x = x ? (M - x) : 0U; return a; }
  ModInt operator+(const ModInt &a) const { return (ModInt(*this) += a); }
  ModInt operator-(const ModInt &a) const { return (ModInt(*this) -= a); }
  ModInt operator*(const ModInt &a) const { return (ModInt(*this) *= a); }
  ModInt operator/(const ModInt &a) const { return (ModInt(*this) /= a); }
  template <class T> friend ModInt operator+(T a, const ModInt &b) { return (ModInt(a) += b); }
  template <class T> friend ModInt operator-(T a, const ModInt &b) { return (ModInt(a) -= b); }
  template <class T> friend ModInt operator*(T a, const ModInt &b) { return (ModInt(a) *= b); }
  template <class T> friend ModInt operator/(T a, const ModInt &b) { return (ModInt(a) /= b); }
  explicit operator bool() const { return x; }
  bool operator==(const ModInt &a) const { return (x == a.x); }
  bool operator!=(const ModInt &a) const { return (x != a.x); }
  friend std::ostream &operator<<(std::ostream &os, const ModInt &a) { return os << a.x; }
};
////////////////////////////////////////////////////////////////////////////////

constexpr unsigned MO = 998244353;
using Mint = ModInt<MO>;


vector<int> uf;
int root(int u) {
  return (uf[u] < 0) ? u : (uf[u] = root(uf[u]));
}
bool connect(int u, int v) {
  u = root(u);
  v = root(v);
  if (u == v) return false;
  if (uf[u] > uf[v]) swap(u, v);
  uf[u] += uf[v];
  uf[v] = u;
  return true;
}


int N, K;
vector<int> A, B;
vector<int> F;

vector<vector<int>> graph;
vector<vector<int>> uss;

int dist[3010][3010];
void dfsDist(int s, int u, int p) {
  for (const int v : graph[u]) if (p != v) {
    dist[s][v] = dist[s][u] + 1;
    dfsDist(s, v, u);
  }
}

vector<vector<int>> GRAPH;
bitset<3010> can[3010];
vector<Mint> dp;

void solve(int k) {
  for (const int l : GRAPH[k]) {
    solve(l);
    for (const int u : uss[k]) {
      Mint sum = 0;
      for (const int v : uss[l]) if (can[u][v]) {
        sum += dp[v];
      }
      dp[u] *= sum;
    }
  }
}

int main() {
  for (int numCases; ~scanf("%d", &numCases); ) { for (int caseId = 1; caseId <= numCases; ++caseId) {
    scanf("%d%d", &N, &K);
    A.resize(N - 1);
    B.resize(N - 1);
    for (int i = 0; i < N - 1; ++i) {
      scanf("%d%d", &A[i], &B[i]);
      --A[i];
      --B[i];
    }
    F.resize(N);
    for (int u = 0; u < N; ++u) {
      scanf("%d", &F[u]);
      --F[u];
    }
    
    graph.assign(N, {});
    for (int i = 0; i < N - 1; ++i) {
      graph[A[i]].push_back(B[i]);
      graph[B[i]].push_back(A[i]);
    }
    
    uss.assign(K, {});
    for (int u = 0; u < N; ++u) uss[F[u]].push_back(u);
    
    uf.assign(N, -1);
    for (int i = 0; i < N - 1; ++i) if (F[A[i]] == F[B[i]]) connect(A[i], B[i]);
    bool ok = true;
    for (int k = 0; k < K; ++k) {
      ok = ok && uss[k].size();
      for (const int u : uss[k]) ok = ok && (root(uss[k][0]) == root(u));
    }
    Mint ans = 0;
    if (ok) {
      for (int s = 0; s < N; ++s) {
        dist[s][s] = 0;
        dfsDist(s, s, -1);
      }
      const int k0 = 0;
      const int rt = uss[k0][0];
      GRAPH.assign(K, {});
      for (int u = 0; u < N; ++u) can[u].reset();
      for (int i = 0; i < N - 1; ++i) if (F[A[i]] != F[B[i]]) {
        int a = A[i], b = B[i];
        if (dist[rt][a] > dist[rt][b]) swap(a, b);
        GRAPH[F[a]].push_back(F[b]);
        for (const int u : uss[F[a]]) for (const int v : uss[F[b]]) {
          can[u][v] = (dist[a][u] == dist[b][v] || dist[a][u] == dist[b][v] + ((F[a] < F[b]) ? +1 : -1));
// if(can[u][v])cerr<<F[a]<<" "<<F[b]<<": "<<u<<" "<<v<<endl;
        }
      }
      dp.assign(N, 1);
      solve(k0);
      for (const int u : uss[k0]) {
        ans += dp[u];
      }
    }
    
    printf("%u\n", ans.x);
  }
#ifndef LOCAL
  break;
#endif
  }
  return 0;
}

Details

Tip: Click on the bar to expand more detailed information

Test #1:

score: 100
Accepted
time: 1ms
memory: 5868kb

input:

10
15 2
10 5
3 5
12 5
10 9
11 7
3 8
2 4
7 1
15 14
8 13
15 6
2 1
4 8
11 15
1 1 1 1 2 1 1 1 2 2 1 2 1 1 1
15 3
8 11
12 8
1 3
13 15
5 9
10 13
6 12
14 4
4 9
15 5
11 10
2 14
7 2
6 3
3 2 3 2 2 3 2 1 2 1 1 3 1 2 1
15 5
1 7
5 2
11 9
6 8
13 3
14 12
3 1
8 9
5 10
10 11
5 1
12 13
10 15
11 4
3 3 3 2 3 2 1 2 2 2 ...

output:

11
15
8
9
5
2
1
0
15
18

result:

ok 10 numbers

Test #2:

score: 0
Accepted
time: 0ms
memory: 5944kb

input:

10
15 2
7 10
10 3
11 4
9 7
4 15
14 8
7 1
8 15
13 4
2 1
15 9
11 6
2 12
5 6
2 2 2 1 1 1 2 1 2 2 1 2 1 1 1
15 3
4 10
12 3
7 3
9 1
6 7
9 11
11 7
10 8
13 15
14 2
14 7
5 15
12 13
11 10
1 2 3 3 3 3 3 3 1 3 3 3 3 2 3
15 5
12 9
6 8
13 2
6 10
2 3
14 15
10 5
4 15
14 13
5 7
11 10
14 12
1 5
13 8
2 4 1 3 2 2 5 3 ...

output:

21
4
5
2
13
17
6
0
0
0

result:

ok 10 numbers

Test #3:

score: 0
Accepted
time: 1593ms
memory: 40824kb

input:

10
3000 2
2052 2631
688 2422
2401 352
654 1669
1566 2157
1187 334
179 178
948 1821
1084 99
878 793
410 336
2218 865
2558 236
2808 430
799 1238
1468 226
312 268
1860 991
2946 96
2540 241
2242 1103
299 1527
788 2026
1885 2104
229 1627
2461 2649
498 2642
1354 98
113 980
947 2790
1281 2232
2682 713
1841...

output:

2420
21248
731916691
864589674
3936
9543168
561554568
0
0
0

result:

ok 10 numbers

Test #4:

score: 0
Accepted
time: 1574ms
memory: 40820kb

input:

10
3000 2
1962 622
2353 451
2284 659
848 1392
2032 111
2647 1540
706 1072
2158 461
887 1075
178 481
2350 2658
2309 96
753 1572
1703 2171
365 1286
2540 2487
2687 194
1585 1241
2147 1247
1720 701
507 2127
2207 812
1568 2621
725 1342
1385 245
2819 2731
2357 450
467 672
2549 1243
1517 702
1808 2415
76 8...

output:

2252
21504
217815511
420767421
3008
36864
187077271
0
0
0

result:

ok 10 numbers

Test #5:

score: 0
Accepted
time: 1036ms
memory: 40852kb

input:

10
3000 2
774 1607
49 867
1384 2838
247 454
1369 2077
2145 1433
1057 1845
1940 2195
1869 967
2830 2829
739 2783
1474 1047
1688 2054
2958 1044
1041 135
768 1621
667 257
1309 1336
2456 2605
2402 1425
1572 1275
2327 1438
1111 1265
607 2583
2557 2788
2694 1378
2332 1848
107 1749
1804 236
2585 309
2761 4...

output:

383799
552301
9347
280893
298918
39575
180102
18385
0
78

result:

ok 10 numbers

Test #6:

score: 0
Accepted
time: 1050ms
memory: 40532kb

input:

10
3000 2
630 377
2445 1749
374 616
1421 937
429 724
2962 2139
1501 2291
2916 2433
2404 2502
755 1952
2756 692
547 420
507 2438
1082 2413
2929 1624
2556 2310
2895 266
2697 2184
1597 502
524 2083
898 1989
190 1269
2189 2104
1751 1387
2986 2041
646 493
627 1899
2516 2434
456 2016
223 1222
2329 2784
94...

output:

7096
400826
74841
21895
331464
110972
215781
220870
0
57

result:

ok 10 numbers

Test #7:

score: 0
Accepted
time: 973ms
memory: 40672kb

input:

10
3000 2
1490 1252
1922 2810
462 631
2806 1842
1552 1043
493 2043
339 2377
1271 354
1043 1795
192 1242
1758 2799
433 452
2473 2444
292 1404
1171 770
1194 2617
2457 2502
124 2059
818 2883
81 835
300 1498
1265 601
2038 1830
1163 1973
1110 1369
1428 392
2061 1037
559 1133
1296 140
2216 565
1001 2432
7...

output:

3074
962904591
573541113
831594950
4847397
673682640
122362172
0
0
0

result:

ok 10 numbers

Test #8:

score: 0
Accepted
time: 1052ms
memory: 40964kb

input:

10
3000 2
459 801
2023 2946
2356 57
2998 264
1498 2840
1755 1603
692 1624
1929 1474
1376 1542
2953 2021
934 2740
1747 1918
270 2928
925 2366
886 2966
1692 2224
1528 1835
1748 2417
76 2539
956 1834
2570 2423
418 146
559 1376
1278 2981
564 555
2200 1472
1006 749
1262 920
2410 429
2149 1023
1365 2864
2...

output:

101017
297761161
615903767
469063443
5029034
178745072
0
0
0
0

result:

ok 10 numbers

Test #9:

score: 0
Accepted
time: 988ms
memory: 40704kb

input:

10
3000 2
1130 2866
2675 1317
1817 1398
1916 550
777 1510
955 1371
2842 2940
1069 86
2086 1600
2635 2465
940 2644
954 2270
2517 895
2981 110
2477 993
2854 2889
247 1409
1008 2137
2436 2896
2588 1174
2774 1946
265 622
1903 100
1052 2594
1864 2791
453 1319
11 2630
2468 1464
952 325
632 2214
1218 2459
...

output:

32539
153019389
251875443
633091978
10314829
918558837
687532775
0
0
0

result:

ok 10 numbers

Test #10:

score: 0
Accepted
time: 979ms
memory: 40660kb

input:

10
3000 2
758 2015
2186 297
1518 1181
2657 1064
687 2929
1271 265
646 408
2332 217
2486 1313
362 1555
2701 1097
2014 2885
1691 43
2804 1610
1866 119
516 740
343 2335
1608 1278
330 1421
1967 2298
94 2157
1392 458
64 216
1090 541
1670 459
1870 750
1780 1973
2276 1910
1744 1699
2602 2901
2484 1380
1422...

output:

159905
198653564
859658661
608555088
9674254
170770457
138273469
0
0
0

result:

ok 10 numbers