/
rerooting.test.cpp
75 lines (68 loc) · 1.5 KB
/
rerooting.test.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#define PROBLEM "https://judge.yosupo.jp/problem/aplusb"
//
#include "../../template/template.hpp"
//
#include "../../graph/graph-template.hpp"
//
#include "../../tree/rerooting.hpp"
//
#include "../../misc/rng.hpp"
#include "../../tree/pruefer-code.hpp"
using namespace Nyaan;
#include "../../modint/montgomery-modint.hpp"
#include "../../modulo/binomial.hpp"
//
using namespace Nyaan;
using mint = LazyMontgomeryModInt<998244353>;
// using mint = LazyMontgomeryModInt<1000000007>;
using vm = vector<mint>;
using vvm = vector<vm>;
Binomial<mint> C;
vm naive(int N, vvi g, vm A) {
auto dfs = [&](auto rc, int c, int p) -> mint {
vm chds;
each(d, g[c]) {
if (d == p) continue;
chds.push_back(rc(rc, d, c));
}
if (chds.empty()) return A[c];
mint prod = 1;
each(x, chds) prod *= x;
return prod + A[c];
};
vm ans(N);
rep(i, N) ans[i] = dfs(dfs, i, -1);
return ans;
}
vm calc(int N, vvi g, vm A) {
mint I = 0;
auto f1 = [&](mint x, mint y) { return x * y; };
auto f2 = [&](mint x, int c, int) { return x + A[c]; };
Rerooting rt{g, f1, f2, I};
auto ans = rt.dp;
rep(i, N) ans[i] += A[i];
return ans;
}
void test() {
rep(t, 10000) {
int N = rng(1, 10);
vvi g = random_tree(N);
vm A(N);
each(x, A) x = rng();
auto an = naive(N, g, A);
auto ac = calc(N, g, A);
assert(an == ac);
}
trc2("OK");
}
void q() {
test();
int a, b;
cin >> a >> b;
cout << a + b << endl;
}
void Nyaan::solve() {
int t = 1;
// in(t);
while (t--) q();
}