QOJ.ac
QOJ
ID | Problem | Submitter | Result | Time | Memory | Language | File size | Submit time | Judge time |
---|---|---|---|---|---|---|---|---|---|
#386226 | #8570. Idola-Tree | hos_lyric | WA | 0ms | 4096kb | C++14 | 9.0kb | 2024-04-11 14:17:06 | 2024-04-11 14:17:06 |
Judging History
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>;
// (without edge property)
// sub[u]: inside subtree at u, rooted at u
// bus[u]: outside subtree at u, rooted at par[u]
// tot[u]: rooted at u
// 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) should attach vertex u to the product of the subtrees.
template <class T> struct ForestDP {
int n;
vector<vector<int>> graph;
vector<int> par;
vector<T> sub, bus, tot;
ForestDP() : n(0) {}
explicit ForestDP(int n_) : n(n_), graph(n_) {}
void ae(int u, int v) {
assert(0 <= u); assert(u < n);
assert(0 <= v); assert(v < n);
graph[u].push_back(v);
graph[v].push_back(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;
for (int j = deg; --j >= 0; ) {
const int v = graph[u][j];
if (p != v) {
dfs0(v, u);
if (~w) {
bus[v].pull(sub[v], bus[w]);
} else {
bus[v] = sub[v];
}
w = v;
}
}
if (~w) {
sub[u] = bus[w];
} else {
sub[u].init();
}
sub[u].up(u);
}
void dfs1(int u, int p) {
const int deg = graph[u].size();
int v = -1;
for (int j = 0; j < deg; ++j) {
const int w = graph[u][j];
if (p != w) {
if (~v) {
bus[v].pull(tot[v], bus[w]);
bus[v].up(u);
tot[w].pull(tot[v], sub[v]);
dfs1(v, u);
} else {
if (~p) {
tot[w] = bus[u];
} else {
tot[w].init();
}
}
v = w;
}
}
if (~v) {
bus[v] = tot[v];
bus[v].up(u);
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);
}
};
/*
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;
Data() {}
void init() {
s0 = s1 = 0;
s2 = 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;
}
void up(int u) {
s2 += s1;
s1 += s0;
s0 += 1;
}
};
// 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();
ForestDP<Data> f(N), g(N);
for (int i = 0; i < N - 1; ++i) {
f.ae(A[i], B[i]);
if (deg[A[i]] > 1 && deg[B[i]] > 1) {
g.ae(A[i], B[i]);
}
}
f.run();
g.run();
// for(int u=0;u<N;++u)cerr<<u<<": "<<f.tot[u].s0<<" "<<f.tot[u].s1<<" "<<f.tot[u].s2<<"; "<<g.tot[u].s0<<" "<<g.tot[u].s1<<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];
ps[i] = 2 * g.tot[y].s1;
}
// 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: 0ms
memory: 3916kb
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: -100
Wrong Answer
time: 0ms
memory: 4096kb
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 230488
result:
wrong answer 4th words differ - expected: '249984', found: '230488'