Wrong answer on test 4.
There is (probably) a much better solution, but I can't figure out what is wrong with this. (only main part of the code)
int n, q, dfn[200001], low[200001], a[200001],
w[200001], an[21][200001], dep[200001], tot = 0, rt = 1;
vector<int> gv[200001];
inline void add_edge(int u, int v){
gv[u].push_back(v);
gv[v].push_back(u);
}
void dfs(int u, int fa){
an[0][u] = fa;
dfn[u] = low[u] = ++tot;
for(int i = 1; i <= 20; i++)
an[i][u] = an[i - 1][an[i - 1][u]];
for(auto v : gv[u]){
if(v == fa) continue;
dep[v] = dep[u] + 1;
dfs(v, u);
low[u] = max(low[u], low[v]);
}
}
int lca(int u, int v){
if(dep[u] < dep[v]) swap(u, v);
for(int i = 20; i >= 0; i--)
if(dep[an[i][u]] >= dep[v]) u = an[i][u];
if(u == v) return u;
for(int i = 20; i >= 0; i--)
if(an[i][u] != an[i][v])
u = an[i][u], v = an[i][v];
return an[0][u];
}
int approach(int u, int v){
if(dep[u] < dep[v]) swap(u, v);
for(int i = 20; i >= 0; i--)
if(dep[an[i][u]] > dep[v]) u = an[i][u];
return u;
}
// SGT
struct node{
int l, r, val, lazy;
node(){
l = 0, r = 0, val = 0, lazy = 0;
}
node operator+(const node &b) const{
node c; c.l = l, c.r = b.r;
c.val = val + b.val;
c.lazy = 0;
return c;
}
}tree[200001 << 2];
void build(int l, int r, int p){
tree[p].l = l, tree[p].r = r;
if(l == r){
tree[p].val = a[l];
tree[p].lazy = 0;
return;
}
int mid = (l + r) / 2;
build(l, mid, p * 2);
build(mid + 1, r, p * 2 + 1);
tree[p] = tree[p * 2] + tree[p * 2 + 1];
}
void mark(int p, int x){
tree[p].val += x * (tree[p].r - tree[p].l + 1);
tree[p].lazy += x;
}
void tag(int p){
if(tree[p].lazy != 0)
mark(p * 2, tree[p].lazy), mark(p * 2 + 1, tree[p].lazy);
tree[p].lazy = 0;
}
void add(int l, int r, int x, int p){
if(l > r) return;
if(l <= tree[p].l && tree[p].r <= r){
mark(p, x); return;
}
tag(p);
int mid = (tree[p].l + tree[p].r) / 2;
if(l <= mid) add(l, r, x, p * 2);
if(mid < r) add(l, r, x, p * 2 + 1);
tree[p] = tree[p * 2] + tree[p * 2 + 1];
}
int query(int l, int r, int p){
if(l > r) return 0;
if(l <= tree[p].l && tree[p].r <= r)
return tree[p].val;
tag(p);
int mid = (tree[p].l + tree[p].r) / 2;
if(l <= mid && mid < r)
return query(l, r, p * 2) + query(l, r, p * 2 + 1);
else if(l <= mid)
return query(l, r, p * 2);
else
return query(l, r, p * 2 + 1);
}
void init_vars(){
// type your initiating code...
}
void solve(int testcase, ...){
init_vars();
cin >> n >> q;
for(int i = 1; i <= n; i++)
cin >> w[i];
for(int i = 1; i < n; i++){
int u, v; cin >> u >> v;
add_edge(u, v);
}
dfs(1, 1);
for(int i = 1; i <= n; i++)
a[dfn[i]] = w[i];
build(1, n, 1);
for(int i = 1; i <= q; i++){
int op, u, v, x; cin >> op;
if(op == 1) cin >> u, rt = u;
else if(op == 2){
cin >> u >> v >> x;
if(lca(u, v) == rt || (lca(u, rt) == rt) + (lca(v, rt) == rt) == 1 || u == rt || v == rt)
add(1, n, x, 1);
else if(lca(u, rt) != lca(u, v) && lca(u, rt) != u && lca(u, rt) != rt && lca(lca(u, v), rt) == lca(u, v))
add(1, dfn[approach(lca(u, rt), rt)] - 1, x, 1), add(low[approach(lca(u, rt), rt)] + 1, n, x, 1);
else if(lca(v, rt) != lca(u, v) && lca(v, rt) != v && lca(v, rt) != rt && lca(lca(u, v), rt) == lca(u, v))
add(1, dfn[approach(lca(v, rt), rt)] - 1, x, 1), add(low[approach(lca(v, rt), rt)] + 1, n, x, 1);
else if(lca(u, rt) != u && lca(u, rt) != rt && lca(v, rt) != v && lca(v, rt) != rt && lca(u, v) != u && lca(u, v) != v)
add(1, dfn[approach(lca(u, v), rt)] - 1, x, 1), add(low[approach(lca(u, v), rt)] + 1, n, x, 1);
else if((lca(u, rt) == u) + (lca(v, rt) == v) == 1){
if(lca(u, rt) == u) add(1, dfn[approach(u, rt)] - 1, x, 1), add(low[approach(u, rt)] + 1, n, x, 1);
else add(1, dfn[approach(v, rt)] - 1, x, 1), add(low[approach(v, rt)] + 1, n, x, 1);
}
else if(lca(u, rt) == u && lca(v, rt) == v){
if(lca(u, v) == u) add(1, dfn[approach(rt, v)] - 1, x, 1), add(low[approach(rt, v)] + 1, n, x, 1);
else add(1, dfn[approach(rt, u)] - 1, x, 1), add(low[approach(rt, u)] + 1, n, x, 1);
}
else
add(dfn[lca(u, v)], low[lca(u, v)], x, 1);
}
else{
cin >> v;
if(v == rt) cout << query(1, n, 1) << endl;
else if(lca(v, rt) != v) cout << query(dfn[v], low[v], 1) << endl;
else cout << query(1, dfn[approach(v, rt)] - 1, 1) + query(low[approach(v, rt)] + 1, n, 1) << endl;
}
}
}