/
subtree_hash.hpp
48 lines (41 loc) · 1.32 KB
/
subtree_hash.hpp
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
#include "mod/modint61.hpp"
#include "graph/base.hpp"
#include "graph/tree.hpp"
#include "random/base.hpp"
#include "graph/tree_dp/rerooting_dp.hpp"
// 複数の木で使って大丈夫
template <typename TREE>
struct SubTree_Hash {
using mint = modint61;
TREE& tree;
vc<u64> dp, dp_1, dp_2;
SubTree_Hash(TREE& tree) : tree(tree) {
int N = tree.N;
using T = pair<int, mint>;
T unit = {0, mint(1)};
auto f_ee = [&](T A, T B) -> T { return {max(A.fi, B.fi), A.se * B.se}; };
auto f_ev = [&](T A, int v) -> T { return {A.fi + 1, A.se}; };
auto f_ve = [&](T A, const auto& e) -> T {
return {A.fi, A.se + hash_base(A.fi)};
};
Rerooting_dp<TREE, T> DP(tree, f_ee, f_ev, f_ve, unit);
dp.resize(N), dp_1.resize(N), dp_2.resize(N);
FOR(v, N) dp[v] = DP.dp[v].se.val;
FOR(v, N) dp_1[v] = DP.dp_1[v].se.val;
FOR(v, N) dp_2[v] = DP.dp_2[v].se.val;
}
// v を根としたときの full tree
u64 operator[](int v) { return dp[v]; }
// root を根としたときの部分木 v
u64 get(int v, int root) {
if (root == v) return dp[v];
if (!tree.in_subtree(root, v)) { return dp_1[v]; }
int w = tree.jump(v, root, 1);
return dp_2[w];
}
static mint hash_base(int k) {
static vc<mint> dat;
while (len(dat) <= k) dat.eb(RNG(mint::get_mod()));
return dat[k];
}
};