QOJ.ac

QOJ

IDProblemSubmitterResultTimeMemoryLanguageFile sizeSubmit timeJudge time
#386245#8570. Idola-Treehos_lyricTL 1ms3800kbC++149.6kb2024-04-11 14:31:462024-04-11 14:31:46

Judging History

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

  • [2024-04-11 14:31:46]
  • 评测
  • 测评结果:TL
  • 用时:1ms
  • 内存:3800kb
  • [2024-04-11 14:31:46]
  • 提交

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>;


// sub[u]: inside subtree at u, rooted at u
// bus[u]: outside subtree at u, rooted at par[u]
// tot[u]: rooted at u
// Edge: edge information
// T: monoid representing information of a subtree.
//   T::init()  should assign the identity.
//   T::pull(const T &l, const T &r)  should assign the product.
//   T::up(int u, int p, const Edge &e)
//       should attach vertex u to the product of the subtrees,
//       and if p != -1 then attach edge information e: u -> p.
template <class Edge, class T> struct ForestDPE {
  int n;
  vector<vector<pair<Edge, int>>> graph;
  vector<int> par;
  vector<T> sub, bus, tot;

  ForestDPE() : n(0) {}
  explicit ForestDPE(int n_) : n(n_), graph(n_) {}
  void ae(int u, int v, const Edge &e0, const Edge &e1) {
    assert(0 <= u); assert(u < n);
    assert(0 <= v); assert(v < n);
    graph[u].emplace_back(e0, v);
    graph[v].emplace_back(e1, u);
  }
  void run() {
    par.assign(n, -2);
    sub.resize(n);
    bus.resize(n);
    tot.resize(n);
    for (int u = 0; u < n; ++u) if (par[u] == -2) {
      dfs0(u, -1);
      dfs1(u, -1);
    }
  }
  void dfs0(int u, int p) {
    par[u] = p;
    const int deg = graph[u].size();
    int w = -1;
    int jp = -1;
    for (int j = deg; --j >= 0; ) {
      const int v = graph[u][j].second;
      if (p != v) {
        dfs0(v, u);
        if (~w) {
          bus[v].pull(sub[v], bus[w]);
        } else {
          bus[v] = sub[v];
        }
        w = v;
      } else {
        jp = j;
      }
    }
    if (~w) {
      sub[u] = bus[w];
    } else {
      sub[u].init();
    }
    sub[u].up(u, p, (~jp) ? graph[u][jp].first : Edge());
  }
  void dfs1(int u, int p) {
    const int deg = graph[u].size();
    int v = -1, jv = -1;
    for (int j = 0; j < deg; ++j) {
      const int w = graph[u][j].second;
      if (p != w) {
        if (~v) {
          bus[v].pull(tot[v], bus[w]);
          bus[v].up(u, v, graph[u][jv].first);
          tot[w].pull(tot[v], sub[v]);
          dfs1(v, u);
        } else {
          if (~p) {
            tot[w] = bus[u];
          } else {
            tot[w].init();
          }
        }
        v = w; jv = j;
      }
    }
    if (~v) {
      bus[v] = tot[v];
      bus[v].up(u, v, graph[u][jv].first);
      tot[u].pull(tot[v], sub[v]);
      dfs1(v, u);
    } else {
      if (~p) {
        tot[u] = bus[u];
      } else {
        tot[u].init();
      }
    }
    tot[u].up(u, -1, Edge());
  }
};

////////////////////////////////////////////////////////////////////////////////


/*
  leaf-edge: w[i]: X[i]-Y[i]
  other: 1
  
  contributions
  1^2: const
  w[i]^2: (N-1)
  1 * 1: const
  w[i] * 1: depends on i
  w[i] * w[j]: 1
    2 \sum[i<j] w[i] w[j] = (\sum[i] w[i])^2 - \sum[i] w[i]^2
*/

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

vector<int> deg;

struct Data {
  Int s0, s1;
  Mint s2;
  // without leaf-edge
  Int t0, t1, dt1;
  
  Data() {}
  void init() {
    s0 = s1 = 0;
    s2 = 0;
    t0 = t1 = dt1 = 0;
  }
  void pull(const Data &a, const Data &b) {
    assert(this != &a);
    assert(this != &b);
    s0 = a.s0 + b.s0;
    s1 = a.s1 + b.s1;
    s2 = a.s2 + b.s2;
    t0 = a.t0 + b.t0;
    t1 = a.t1 + b.t1;
    dt1 = a.dt1 + b.dt1;
  }
  void up(int u, int p, int e) {
    s2 += s1;
    s1 += s0;
    s0 += 1;
    t1 += dt1;
    t0 += 1;
    dt1 = e * t0;
  }
};

// merge (a[i], a[i] + d, a[i] + 2 d, ...)
template <class F> void mergeArithBrute(const vector<Int> &as, Int d, int n, F f) {
  const int asLen = as.size();
  priority_queue<Int, vector<Int>, greater<Int>> que;
  for (int i = 0; i < asLen; ++i) que.push(as[i]);
  for (int j = 0; j < n; ++j) {
    const Int b = que.top();
    que.pop();
    f(b);
    que.push(b + d);
  }
}

int main() {
  for (int numCases; ~scanf("%d", &numCases); ) { for (int caseId = 1; caseId <= numCases; ++caseId) {
    scanf("%d%d", &N, &C);
    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];
    }
    
    Mint ans = 0;
    if (N == 2) {
      for (int c = N - 1; c <= C; ++c) {
        const Mint now = (Int)c * c;
        ans += now*now*now;
      }
    } else {
      deg.assign(N, 0);
      for (int i = 0; i < N - 1; ++i) {
        ++deg[A[i]];
        ++deg[B[i]];
      }
      int rt = -1;
      vector<int> X;
      for (int u = 0; u < N; ++u) {
        if (deg[u] == 1) {
          X.push_back(u);
        } else {
          if (!~rt) rt = u;
        }
      }
      const int L = X.size();
      
      ForestDPE<int, Data> f(N);
      for (int i = 0; i < N - 1; ++i) {
        const int e = (deg[A[i]] > 1 && deg[B[i]] > 1) ? 1 : 0;
        f.ae(A[i], B[i], e, e);
      }
      f.run();
// for(int u=0;u<N;++u)cerr<<u<<": "<<f.tot[u].s0<<" "<<f.tot[u].s1<<" "<<f.tot[u].s2<<"; "<<f.tot[u].t0<<" "<<f.tot[u].t1<<endl;
      
      // all 1
      Mint base = 0;
      for (int u = 0; u < N; ++u) {
        base += f.tot[u].s1;
        base += 2 * f.tot[u].s2;
      }
      base /= 2;
      
      // (N-2) w^2 + p w
      vector<Int> ps(L);
      for (int i = 0; i < L; ++i) {
        const int x = X[i];
        const int y = f.graph[x][0].second;
        ps[i] = 2 * f.tot[y].t1;
      }
// cerr<<"base = "<<base<<", ps = "<<ps<<endl;
      
      // diff: (N-2) (2w+1) + p
      for (int i = 0; i < L; ++i) {
        ps[i] += (N-2) * 3;
      }
      
      Mint now = base;
      int c = L;
      mergeArithBrute(ps, 2 * (N-2), C - (N-1) + 1, [&](Int t) -> void {
// cerr<<"c = "<<c<<": now = "<<now<<", t = "<<t<<endl;
        ans += now*now*now;
        now += t;
        // c -> c+1
        now += (c + c + 1);
        ++c;
      });
    }
    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: 3800kb

input:

2
4 3
1 4
1 3
2 1
4 4
1 4
1 3
2 1

output:

3375
25327

result:

ok 2 tokens

Test #2:

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

input:

4
4 3
1 4
1 3
2 1
4 4
1 4
1 3
2 1
5 4
1 4
1 3
1 2
5 4
5 5
1 4
1 3
1 2
5 4

output:

3375
25327
54872
249984

result:

ok 4 tokens

Test #3:

score: -100
Time Limit Exceeded

input:

4
300000 50000000
216838 200677
44849 12926
125086 157230
26195 29767
241694 21336
21336 24457
105861 84565
184655 45583
175336 97945
286044 30927
295273 249694
109469 1566
193560 251229
176229 288707
206166 13532
166218 264413
299977 95587
159105 48116
57912 82606
97296 193689
115029 121192
68068 1...

output:


result: