牛客Wannafly挑战赛26E 蚂蚁开会(树链剖分+线段树)

传送门

题面描述

一颗n个节点的树，m次操作，有点权（该节点蚂蚁个数）和边权（相邻节点的距离）。

三种操作：

操作1：1 i x将节点i的点权修改为x。(1 <= i <= n; 1 <= x <= 100000)

操作2：2 i x将第i条边的边权修改为x。(1 <= i < n; 1 <= x <= 100000)

操作3：3 i 节点i发出开会指令，求树上所有蚂蚁到走到节点i的距离和。(1 <= i <= n)

题解

先转换问题为

求:

\(\sum_{i = 1}^{n} dep[x] + dep[i] - 2 * dep[lca(x, i)]\)

前面两个式子显然可以\(O(1)\)求

问题就是求后面那个东西，

可以发现,要求的就是两倍的\(1-i,1-x\)重复的部分

那么我们树链剖分一下，记下每条边被\(1-i\)这样的路径经过了多少次

然后询问x，就是求\(1-x\)这条路径上每一条边的 边权 * 被经过次数

接着就是想办法维护上面的东西

线段树维护两个东西,区间边权和 \(c\),边权 * 被经过次数和 \(t\)

操作\(1:(x,y)\)

设原来权值为\(a[x]\)

修改即是,给点\(x\)加上\((y-a[x])\)

操作\(2:(i, x)\)

对于第一个式子,\(\sum_{i = 1}^{n}dep[x]\)

我们修改了这条边,那么以\(i\)为根子树所有点\(dep\)都要减

所以要维护一棵树状数组,来统计子树有多少只蚂蚁

然后修改线段树,

因为是单点修改,递归找到那条边

直接修改后,回溯更新即可

Code

#include<bits/stdc++.h>

#define int long long
#define RG register

using namespace std;
template<class T> inline void read(T &x) {
	x = 0; RG char c = getchar(); bool f = 0;
	while (c != '-' && (c < '0' || c > '9')) c = getchar(); if (c == '-') c = getchar(), f = 1;
	while (c >= '0' && c <= '9') x = x*10+c-48, c = getchar();
	x = f ? -x : x;
	return ;
}
template<class T> inline void write(T x) {
	if (!x) {putchar(48);return ;}
	if (x < 0) x = -x, putchar('-');
	int len = -1, z[20]; while (x > 0) z[++len] = x%10, x /= 10;
	for (RG int i = len; i >= 0; i--) putchar(z[i]+48);return ;
}

const int N = 500010;

struct node {
	int to, nxt, w;
}g[N];
int last[N], gl, n, m, a[N];
void add(int x, int y, int z) {
	g[++gl] = (node) {y, last[x], z};
	last[x] = gl;
}

int siz[N], son[N], s[N], dfn[N], w[N], tim, top[N], ed[N], fa[N];

void dfs1(int u) {
	siz[u] = 1; int mx = 0;
	for (int i = last[u]; i; i = g[i].nxt) {
		int v = g[i].to;
		fa[v] = u; s[v] = g[i].w;
		dfs1(v); siz[u] += siz[v];
		if (mx < siz[v]) mx = siz[v], son[u] = v;
	}
	return ;
}

void dfs2(int u, int topf) {	
	top[u] = topf; w[++tim] = s[u]; dfn[u] = tim;
	if (!son[u]) {ed[u] = tim; return ;}
	dfs2(son[u], topf);
	for (int i = last[u]; i; i = g[i].nxt) {
		int v = g[i].to; if (v == son[u]) continue;
		dfs2(v, v);		
	}
	ed[u] = tim;
	return ;
}

int c[N << 2], t[N << 2], lz[N << 2], sum, gg, ans;
#define ls (rt << 1)
#define rs (rt << 1 | 1)
#define mid ((l + r) >> 1)
void pushdown(int rt) {lz[ls] += lz[rt]; lz[rs] += lz[rt];t[ls] += lz[rt] * c[ls]; t[rs] += lz[rt] * c[rs];lz[rt] = 0;}
void pushup(int rt) {t[rt] = t[ls] + t[rs];c[rt] = c[ls] + c[rs];}
void build(int rt, int l, int r) {
	if (l == r) {c[rt] = w[l]; return ;}
	build(ls, l, mid); build(rs, mid + 1, r);
	pushup(rt);
}
void update(int rt, int l, int r, int L, int R, int k) {
	if (L <= l && r <= R) {
		lz[rt] += k; t[rt] += k * c[rt]; sum += k * c[rt];
		return ;
	}
	if (lz[rt]) pushdown(rt);
	if (L <= mid) update(ls, l, mid, L, R, k);
	if (R > mid) update(rs, mid + 1, r, L, R, k);
	pushup(rt);
}

void update(int rt, int l, int r, int pos, int k) {
	if (l == r) {
		t[rt] = t[rt] / c[rt] * k, c[rt] = k;
		return ;
	}
	if (lz[rt]) pushdown(rt);
	if (pos <= mid) update(ls, l, mid, pos, k);
	else update(rs, mid + 1, r, pos, k);
	pushup(rt);
}


int T[N];
inline void add(int x, int k) {for (; x <= n; x += x & (-x)) T[x] += k;}
inline int get(int x) {int res = 0;for (; x; x -= x & (-x)) res += T[x];return res;}

int query(int rt, int l, int r, int L, int R) {
	if (L <= l && r <= R) {sum += c[rt];return t[rt];}
	int res = 0;
	if (lz[rt]) pushdown(rt);
	if (L <= mid) res = query(ls, l, mid, L, R);
	if (R > mid) res += query(rs, mid + 1, r, L, R);
	return res;
}

int query(int x) {
	sum = 0;
	int res = 0;
	while (x) {
		res += query(1, 1, n, dfn[top[x]], dfn[x]);
		x = fa[top[x]];
	}
	return ans + sum * gg - 2 * res;
}

void update(int x, int k) {
	sum = 0; gg += k;
	add(dfn[x], k);
	while (x) {
		update(1, 1, n, dfn[top[x]], dfn[x], k);
		x = fa[top[x]];
	}
	ans += sum;
	return ;
}

signed main() {
	read(n), read(m);
	for (int i = 1; i <= n; i++) read(a[i]);
	for (int i = 2, y, z; i <= n; i++) {
		read(y), read(z);
		add(y, i, z);
	}
	dfs1(1), dfs2(1, 1);
	build(1, 1, n);
	for (int i = 1; i <= n; i++) update(i, a[i]);
	while (m--) {
		int op, x, y;
		read(op), read(x);
		if (op == 3) printf("%lld\n", query(x));
		else {
			read(y);
			if (op == 1) update(x, y - a[x]), a[x] = y;
			else {
				ans += (y - s[x + 1]) * (get(ed[x + 1]) - get(dfn[x + 1] - 1));
				update(1, 1, n, dfn[x + 1], y), s[x + 1] = y;
			} 
		}
	}
	return 0;
}