QOJ.ac

QOJ

IDProblemSubmitterResultTimeMemoryLanguageFile sizeSubmit timeJudge time
#298634#7895. Graph Partitioning 2ucup-team087#TL 1366ms37156kbC++1410.7kb2024-01-06 13:41:422024-01-06 13:41:43

Judging History

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

  • [2024-01-06 13:41:43]
  • 评测
  • 测评结果:TL
  • 用时:1366ms
  • 内存:37156kb
  • [2024-01-06 13:41: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 = 998244353U;
constexpr unsigned MO2 = 2U * MO;
constexpr int FFT_MAX = 23;
using Mint = ModInt<MO>;
constexpr Mint FFT_ROOTS[FFT_MAX + 1] = {1U, 998244352U, 911660635U, 372528824U, 929031873U, 452798380U, 922799308U, 781712469U, 476477967U, 166035806U, 258648936U, 584193783U, 63912897U, 350007156U, 666702199U, 968855178U, 629671588U, 24514907U, 996173970U, 363395222U, 565042129U, 733596141U, 267099868U, 15311432U};
constexpr Mint INV_FFT_ROOTS[FFT_MAX + 1] = {1U, 998244352U, 86583718U, 509520358U, 337190230U, 87557064U, 609441965U, 135236158U, 304459705U, 685443576U, 381598368U, 335559352U, 129292727U, 358024708U, 814576206U, 708402881U, 283043518U, 3707709U, 121392023U, 704923114U, 950391366U, 428961804U, 382752275U, 469870224U};
constexpr Mint FFT_RATIOS[FFT_MAX] = {911660635U, 509520358U, 369330050U, 332049552U, 983190778U, 123842337U, 238493703U, 975955924U, 603855026U, 856644456U, 131300601U, 842657263U, 730768835U, 942482514U, 806263778U, 151565301U, 510815449U, 503497456U, 743006876U, 741047443U, 56250497U, 867605899U};
constexpr Mint INV_FFT_RATIOS[FFT_MAX] = {86583718U, 372528824U, 373294451U, 645684063U, 112220581U, 692852209U, 155456985U, 797128860U, 90816748U, 860285882U, 927414960U, 354738543U, 109331171U, 293255632U, 535113200U, 308540755U, 121186627U, 608385704U, 438932459U, 359477183U, 824071951U, 103369235U};

// as[rev(i)] <- \sum_j \zeta^(ij) as[j]
void fft(Mint *as, int n) {
  assert(!(n & (n - 1))); assert(1 <= n); assert(n <= 1 << FFT_MAX);
  int m = n;
  if (m >>= 1) {
    for (int i = 0; i < m; ++i) {
      const unsigned x = as[i + m].x;  // < MO
      as[i + m].x = as[i].x + MO - x;  // < 2 MO
      as[i].x += x;  // < 2 MO
    }
  }
  if (m >>= 1) {
    Mint prod = 1U;
    for (int h = 0, i0 = 0; i0 < n; i0 += (m << 1)) {
      for (int i = i0; i < i0 + m; ++i) {
        const unsigned x = (prod * as[i + m]).x;  // < MO
        as[i + m].x = as[i].x + MO - x;  // < 3 MO
        as[i].x += x;  // < 3 MO
      }
      prod *= FFT_RATIOS[__builtin_ctz(++h)];
    }
  }
  for (; m; ) {
    if (m >>= 1) {
      Mint prod = 1U;
      for (int h = 0, i0 = 0; i0 < n; i0 += (m << 1)) {
        for (int i = i0; i < i0 + m; ++i) {
          const unsigned x = (prod * as[i + m]).x;  // < MO
          as[i + m].x = as[i].x + MO - x;  // < 4 MO
          as[i].x += x;  // < 4 MO
        }
        prod *= FFT_RATIOS[__builtin_ctz(++h)];
      }
    }
    if (m >>= 1) {
      Mint prod = 1U;
      for (int h = 0, i0 = 0; i0 < n; i0 += (m << 1)) {
        for (int i = i0; i < i0 + m; ++i) {
          const unsigned x = (prod * as[i + m]).x;  // < MO
          as[i].x = (as[i].x >= MO2) ? (as[i].x - MO2) : as[i].x;  // < 2 MO
          as[i + m].x = as[i].x + MO - x;  // < 3 MO
          as[i].x += x;  // < 3 MO
        }
        prod *= FFT_RATIOS[__builtin_ctz(++h)];
      }
    }
  }
  for (int i = 0; i < n; ++i) {
    as[i].x = (as[i].x >= MO2) ? (as[i].x - MO2) : as[i].x;  // < 2 MO
    as[i].x = (as[i].x >= MO) ? (as[i].x - MO) : as[i].x;  // < MO
  }
}

// as[i] <- (1/n) \sum_j \zeta^(-ij) as[rev(j)]
void invFft(Mint *as, int n) {
  assert(!(n & (n - 1))); assert(1 <= n); assert(n <= 1 << FFT_MAX);
  int m = 1;
  if (m < n >> 1) {
    Mint prod = 1U;
    for (int h = 0, i0 = 0; i0 < n; i0 += (m << 1)) {
      for (int i = i0; i < i0 + m; ++i) {
        const unsigned long long y = as[i].x + MO - as[i + m].x;  // < 2 MO
        as[i].x += as[i + m].x;  // < 2 MO
        as[i + m].x = (prod.x * y) % MO;  // < MO
      }
      prod *= INV_FFT_RATIOS[__builtin_ctz(++h)];
    }
    m <<= 1;
  }
  for (; m < n >> 1; m <<= 1) {
    Mint prod = 1U;
    for (int h = 0, i0 = 0; i0 < n; i0 += (m << 1)) {
      for (int i = i0; i < i0 + (m >> 1); ++i) {
        const unsigned long long y = as[i].x + MO2 - as[i + m].x;  // < 4 MO
        as[i].x += as[i + m].x;  // < 4 MO
        as[i].x = (as[i].x >= MO2) ? (as[i].x - MO2) : as[i].x;  // < 2 MO
        as[i + m].x = (prod.x * y) % MO;  // < MO
      }
      for (int i = i0 + (m >> 1); i < i0 + m; ++i) {
        const unsigned long long y = as[i].x + MO - as[i + m].x;  // < 2 MO
        as[i].x += as[i + m].x;  // < 2 MO
        as[i + m].x = (prod.x * y) % MO;  // < MO
      }
      prod *= INV_FFT_RATIOS[__builtin_ctz(++h)];
    }
  }
  if (m < n) {
    for (int i = 0; i < m; ++i) {
      const unsigned y = as[i].x + MO2 - as[i + m].x;  // < 4 MO
      as[i].x += as[i + m].x;  // < 4 MO
      as[i + m].x = y;  // < 4 MO
    }
  }
  const Mint invN = Mint(n).inv();
  for (int i = 0; i < n; ++i) {
    as[i] *= invN;
  }
}

void fft(vector<Mint> &as) {
  fft(as.data(), as.size());
}
void invFft(vector<Mint> &as) {
  invFft(as.data(), as.size());
}

vector<Mint> convolve(vector<Mint> as, vector<Mint> bs) {
  if (as.empty() || bs.empty()) return {};
  const int len = as.size() + bs.size() - 1;
  int n = 1;
  for (; n < len; n <<= 1) {}
  as.resize(n); fft(as);
  bs.resize(n); fft(bs);
  for (int i = 0; i < n; ++i) as[i] *= bs[i];
  invFft(as);
  as.resize(len);
  return as;
}
vector<Mint> square(vector<Mint> as) {
  if (as.empty()) return {};
  const int len = as.size() + as.size() - 1;
  int n = 1;
  for (; n < len; n <<= 1) {}
  as.resize(n); fft(as);
  for (int i = 0; i < n; ++i) as[i] *= as[i];
  invFft(as);
  as.resize(len);
  return as;
}
////////////////////////////////////////////////////////////////////////////////


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

vector<vector<int>> graph;

struct Result {
  int off;
  vector<Mint> fs;
  Mint operator[](int i) const {
    return (0 <= i - off && i - off < (int)fs.size()) ? fs[i - off] : 0;
  }
};
Result solve(int u, int p) {
  Result crt;
  crt.off = 0;
  crt.fs = {0, 1};
  for (const int v : graph[u]) if (p != v) {
    const Result res = solve(v, u);
    // cut
    const Mint cut = res[K] + res[K + 1];
    Result nxt = crt;
    for (auto &f : nxt.fs) f *= cut;
    // connect
    const auto prod = convolve(crt.fs, res.fs);
    int len = crt.off + res.off + (int)prod.size();
    chmin(len, K + 2);
    if (nxt.off + (int)nxt.fs.size() < len) {
      nxt.fs.resize(len - nxt.off, 0);
    }
    for (int i = 0; i < (int)prod.size(); ++i) {
      if (crt.off + res.off + i < K + 2) {
        nxt.fs[res.off + i] += prod[i];
      }
    }
    crt = nxt;
  }
  return crt;
}

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];
    }
    
    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]);
    }
    const auto res = solve(0, -1);
    const Mint ans = res[K] + res[K + 1];
    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: 3700kb

input:

2
8 2
1 2
3 1
4 6
3 5
2 4
8 5
5 7
4 3
1 2
1 3
2 4

output:

2
1

result:

ok 2 lines

Test #2:

score: 0
Accepted
time: 209ms
memory: 4240kb

input:

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

output:

0
3
112
0
1
0
1
0
0
0
1
0
1
0
0
1
0
140
0
0
0
814
1
6
1
1
2
2
0
612
0
1
0
0
0
1
1
0
0
121
4536
0
0
1718
0
0
1
0
444
1
1908
1813
3
74
0
1
0
46
0
0
0
0
0
0
0
0
0
1
0
1
1
1
239
0
0
0
1
0
0
0
1
0
1
0
0
1
1
0
0
0
1
0
0
0
48
0
2
0
0
0
1
364
0
206
0
0
76
0
1
0
0
2
0
1
2
0
0
1
0
0
4
0
1
1
0
0
1
1
1
0
0
1
1
...

result:

ok 5550 lines

Test #3:

score: 0
Accepted
time: 394ms
memory: 10528kb

input:

3
99990 259
23374 69108
82204 51691
8142 67119
48537 97966
51333 44408
33147 68485
21698 86824
15746 58746
78761 86975
58449 61819
69001 68714
25787 2257
25378 14067
64899 68906
29853 31359
75920 85420
76072 11728
63836 55505
43671 98920
77281 25176
40936 66517
61029 61440
66908 52300
92101 59742
69...

output:

259200
247
207766300

result:

ok 3 lines

Test #4:

score: 0
Accepted
time: 512ms
memory: 12152kb

input:

3
99822 332
11587 83046
63424 60675
63423 73718
74622 40130
5110 26562
28361 80899
30886 70318
8708 11068
34855 96504
7904 75735
31904 42745
87892 55105
82374 81319
77407 82147
91475 12343
13470 95329
58766 95716
83232 44156
75907 92437
69785 93598
47857 33018
62668 31394
24238 72675
98254 43583
180...

output:

315881300
4505040
185631154

result:

ok 3 lines

Test #5:

score: 0
Accepted
time: 717ms
memory: 12168kb

input:

3
99021 1000
41739 4318
72541 76341
31227 15416
49232 13808
50837 51259
74464 11157
92684 84646
95226 64673
74155 82511
33301 31373
5901 29318
38227 98893
96752 57411
35167 42401
24344 90803
6956 33753
51120 24535
29594 2646
70305 32961
93079 38070
49273 48987
62799 77986
94353 84447
74970 31546
263...

output:

917568
1
1213

result:

ok 3 lines

Test #6:

score: 0
Accepted
time: 923ms
memory: 11504kb

input:

3
100000 10000
15556 26349
14984 68012
43040 63434
32168 60646
70509 38559
26238 29031
45952 19431
29510 54395
5676 59515
28220 41438
46902 56640
8221 80059
77225 66558
17473 87324
20819 35098
29515 12641
84108 41157
11223 66562
25999 95852
3790 63605
20963 15799
217 58841
61619 13324
3406 60525
693...

output:

1
1
1

result:

ok 3 lines

Test #7:

score: 0
Accepted
time: 1366ms
memory: 37156kb

input:

3
99969 79
84806 29026
76190 49303
81448 88606
47775 83229
7424 30063
75504 59640
28456 18012
26623 18383
66305 32640
55981 65140
25760 523
76248 13482
25598 55231
96844 17032
44892 64592
4915 50521
49879 86466
99286 20894
97915 76337
38424 2546
17489 46475
91791 2636
79283 35209
14773 60224
94096 5...

output:

855988479
413863362
390147247

result:

ok 3 lines

Test #8:

score: -100
Time Limit Exceeded

input:

3
99655 347
11149 99084
14300 87239
74978 75669
48703 12705
62600 62372
85743 67544
11248 36566
31920 23357
91970 67460
47599 56467
67521 16526
50284 63800
6701 3456
15660 81507
43192 5734
57965 67731
42676 26195
60309 19848
30504 47635
66455 98017
1645 70119
47861 95592
32453 39251
31178 59516
2144...

output:


result: