Heavy-Light Decomposition
Path Queries II
I’m stuck on this problem, and I don’t understand why my solution with O(n * log(n) * log(n)) asymptotics gets TLE and takes twice as long as the author’s solution. Please help me. Here’s my code:
#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
using u64 = unsigned long long;
int log(int n) {
int ans = 0;
while (n > 1) {
ans++;
n >>= 1;
}
return ans;
}
struct Info {
int max = 0;
};
Info operator+(const Info& a, const Info& b) {
return {std::max(a.max, b.max)};
}
template<class Info>
struct SegmentTree {
int n;
vector<Info> info;
SegmentTree() : n(0) {}
SegmentTree(int n_, Info v_ = Info()) {
init(n_, v_);
}
template<class T>
SegmentTree(vector<T> init_) {
init(init_);
}
void init(int n_, Info v_ = Info()) {
init(vector<Info>(n_, v_));
}
template<class T>
void init(vector<T> init_) {
n = init_.size();
info.assign(4 << log(n), Info());
auto build = [&](auto& self, int p, int l, int r) -> void {
if (r - l == 1) {
info[p] = init_[l];
return;
}
int m = (l + r) / 2;
self(self, 2 * p + 1, l, m);
self(self, 2 * p + 2, m, r);
pull(p);
};
build(build, 0, 0, n);
}
void pull(int p) {
info[p] = info[2 * p + 1] + info[2 * p + 2];
}
void modify(int p, int l, int r, int x, const Info& v) {
if (r - l == 1) {
info[p] = v;
return;
}
int m = (l + r) / 2;
if (x < m) {
modify(2 * p + 1, l, m, x, v);
}
else {
modify(2 * p + 2, m, r, x, v);
}
pull(p);
}
void modify(int x, const Info& v) {
modify(0, 0, n, x, v);
}
Info rangeQuery(int p, int l, int r, int x, int y) {
if (l >= y || r <= x) {
return Info();
}
if (l >= x && r <= y) {
return info[p];
}
int m = (l + r) / 2;
return rangeQuery(2 * p + 1, l, m, x, y) + rangeQuery(2 * p + 2, m, r, x, y);
}
Info rangeQuery(int x, int y) {
return rangeQuery(0, 0, n, x, y);
}
};
struct HLD {
int n;
vector<int> siz, top, dep, parent, in, out, seq;
vector<vector<int>> adj;
int cur;
HLD() {}
HLD(int n) {
init(n);
}
void init(int n) {
this->n = n;
siz.resize(n);
top.resize(n);
dep.resize(n);
parent.resize(n);
in.resize(n);
out.resize(n);
seq.resize(n);
adj.assign(n, {});
cur = 0;
}
void addEdge(int u, int v) {
adj[u].push_back(v);
adj[v].push_back(u);
}
void work(int root = 0) {
top[root] = root;
parent[root] = -1;
dfs1(root);
dfs2(root);
}
void dfs1(int u) {
if (parent[u] != -1) {
adj[u].erase(find(adj[u].begin(), adj[u].end(), parent[u]));
}
siz[u] = 1;
for (auto& v : adj[u]) {
parent[v] = u;
dep[v] = dep[u] + 1;
dfs1(v);
siz[u] += siz[v];
if (siz[v] > siz[adj[u][0]]) {
swap(v, adj[u][0]);
}
}
}
void dfs2(int u) {
in[u] = cur++;
seq[in[u]] = u;
for (auto v : adj[u]) {
top[v] = v == adj[u][0] ? top[u] : v;
dfs2(v);
}
out[u] = cur;
}
};
int main() {
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
int n, q;
cin >> n >> q;
vector<int> v(n);
for (int i = 0; i < n; i++) {
cin >> v[i];
}
HLD t(n);
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
u--, v--;
t.addEdge(u, v);
}
t.work();
SegmentTree<Info> seg(n);
for (int i = 0; i < n; i++) {
seg.modify(t.in[i], {v[i]});
}
while (q--) {
int op;
cin >> op;
if (op == 1) {
int s, x;
cin >> s >> x;
s--;
seg.modify(t.in[s], {x});
}
else {
int a, b;
cin >> a >> b;
a--, b--;
int ans = 0;
while (t.top[a] != t.top[b]) {
if (t.dep[t.top[b]] > t.dep[t.top[a]]) {
swap(a, b);
}
ans = std::max(ans, seg.rangeQuery(t.in[t.top[a]], t.in[a] + 1).max);
a = t.parent[t.top[a]];
}
if (t.dep[a] > t.dep[b]) {
swap(a, b);
}
ans = std::max(ans, seg.rangeQuery(t.in[a], t.in[b] + 1).max);
cout << ans << " ";
}
}
}