さあ次はいよいよ HL分解の練習をするのん!!!!!
問題概要
N 頂点のツリーが与えられて以下の Q 個の操作を行う。初期状態では全頂点の値は 0 である。また最終的に出力する値 res の初期値も 0 とする
- 2 頂点 u, v 間のパス上に含まれる頂点すべて (u, v も含む) の値を +1 する。
- その後、パス u-v に含まれる頂点すべて (u, v も含む) の値の合計を求め、それを res に合算する
最後に res を出力せよ。
制約
- 2 <= N, Q <= 105
考えたこと
HL 分解を勉強した。HL分解は、ツリーを
- 高々 本のパスに分解する
- 各パスごとにセグ木なりを載せて処理する
という感じのことをする。おおむね 1 クエリあたり の計算量になる。
HL 分解を頑張る流れは、パスに関する問題であれば
という感じにできるので、実質「列に対するクエリ処理方法を考える」問題に落ちる。Euler ツアーと比較すると「逆演算」を考えなくてよいので汎用性がさらに高そう。
今回の問題で言えば、「区間加算」と「区間和取得」は bit を 2 つ持つ方法 (蟻本に載ってる) で実現することができる。これらをラムダ式で以下のように渡すだけでよい。文字通り貼るだけ感がある。
hld.foreach_nodes(a, b, [&](int l, int r) { bit.add(l+1, r+2, 1); res += bit.sum(l+1, r+2); });
#include <iostream> #include <functional> #include <vector> #include <queue> #include <stack> using namespace std; // BIT template <class Abel> struct BIT { vector<Abel> dat[2]; Abel UNITY_SUM = 0; // to be set /* [1, n] */ BIT(int n) { init(n); } void init(int n) { for (int iter = 0; iter < 2; ++iter) dat[iter].assign(n + 1, UNITY_SUM); } /* a is 1-indexed */ inline void sub_add(int p, int a, Abel x) { for (int i = a; i < (int)dat[p].size(); i += i & -i) dat[p][i] = dat[p][i] + x; } inline void add(int a, int b, Abel x) { sub_add(0, a, x * -(a - 1)); sub_add(1, a, x); sub_add(0, b, x * (b - 1)); sub_add(1, b, x * (-1)); } /* [1, a], a is 1-indexed */ inline Abel sub_sum(int p, int a) { Abel res = UNITY_SUM; for (int i = a; i > 0; i -= i & -i) res = res + dat[p][i]; return res; } inline Abel sum(int a, int b) { return sub_sum(0, b - 1) + sub_sum(1, b - 1) * (b - 1) - sub_sum(0, a - 1) - sub_sum(1, a - 1) * (a - 1); } /* debug */ void print() { for (int i = 1; i < (int)dat[0].size(); ++i) cout << sum(i, i + 1) << ","; cout << endl; } }; // HL-Decomposition // vid: id of v after HL-Decomposition // inv: inv[vid[v]] = v // par: id of parent // depth // subsize: size of subtree // head: head-id in the heavy-path // next: next-id in the heavy-oatg // type: the id of tree for forest typedef vector<vector<int> > Graph; struct HLDecomposition { int n; Graph G; vector<int> vid, inv, par, depth, subsize, head, next, type; // construct HLDecomposition() { } HLDecomposition(const Graph &G_) : n((int)G_.size()), G(G_), vid(n, -1), inv(n), par(n), depth(n), subsize(n, 1), head(n), next(n, -1), type(n) { } void build(vector<int> roots = {0}) { int curtype = 0, pos = 0; for (auto r : roots) dfs(r), bfs(r, curtype++, pos); } void dfs(int r) { stack<pair<int,int> > st; par[r] = -1, depth[r] = 0; st.emplace(r, 0); while (!st.empty()) { int v = st.top().first; int &i = st.top().second; if (i < (int)G[v].size()) { int e = G[v][i++]; if (e == par[v]) continue; par[e] = v, depth[e] = depth[v] + 1; st.emplace(e, 0); } else { st.pop(); int maxsize = 0; for (auto e : G[v]) { if (e == par[v]) continue; subsize[v] += subsize[e]; if (maxsize < subsize[e]) maxsize = subsize[e], next[v] = e; } } } } void bfs(int r, int curtype, int &pos) { queue<int> que({r}); while (!que.empty()) { int start = que.front(); que.pop(); for (int v = start; v != -1; v = next[v]) { type[v] = curtype; vid[v] = pos++; inv[vid[v]] = v; head[v] = start; for (auto e : G[v]) if (e != par[v] && e != next[v]) que.push(e); } } } // node query [u, v], f([left, right]) void foreach_nodes(int u, int v, const function<void(int,int)> &f) { while (true) { if (vid[u] > vid[v]) swap(u, v); f(max(vid[head[v]], vid[u]), vid[v]); if (head[u] != head[v]) v = par[head[v]]; else break; } } // edge query [u, v], f([left, right]) void foreach_edges(int u, int v, const function<void(int,int)> &f) { while (true) { if (vid[u] > vid[v]) swap(u, v); if (head[u] != head[v]) { f(vid[head[v]], vid[v]); v = par[head[v]]; } else { if (u != v) f(vid[u] + 1, vid[v]); break; } } } // LCA int lca(int u, int v) { while (true) { if (vid[u] > vid[v]) swap(u, v); if (head[u] == head[v]) return u; v = par[head[v]]; } } }; int main() { int N, Q; cin >> N; Graph G(N); for (int i = 0; i < N-1; ++i) { int u, v; scanf("%d %d", &u, &v); --u, --v; G[u].push_back(v); G[v].push_back(u); } HLDecomposition hld(G); hld.build(); long long res = 0; BIT<long long> bit(N+1); cin >> Q; for (int q = 0; q < Q; ++q) { int a, b; scanf("%d %d", &a, &b); --a, --b; hld.foreach_nodes(a, b, [&](int l, int r) { bit.add(l+1, r+2, 1); res += bit.sum(l+1, r+2); }); } cout << res << endl; }