https://codeforces.com/contest/916/problem/E
有一棵n个点的树,每个节点上有一个权值wi,最开始根为1号点.现在有3种
类型的操作:
• 1 root, 表示将根设为root.
• 2 u v x, 设u, v的最近公共祖先为p, 将p的子树中的所有点的权值加上x.
• 3 u, 查询u的子树中的所有点的权值和.
对于类型3操作,输出答案.
线段树 + 树链剖分
树链剖分的意义在于:由于函数
int lca(int x,int y)
大量使用,如果不使用树链剖分,当树退化为一根链时,
int lca(int x,int y)的复杂度为O(n), 从而整个算法的复杂度为O(n^2)
使用树链剖分时,整个算法的复杂度为O(nlogn)
该程序的关键点是:
不实施换根,每次只记录新根,并注意新根在后续操作的影响。
主要难点:在不同新根位置下,求包含x,y 的最小子树的根节点
#include "stdafx.h"
#include<cstdio>
#include<cstring>
#include<iostream>
#include<cmath>
#include<cstdlib>
#include<algorithm>
using namespace std;
typedef long long LL;
const int N = 1e5 + 10,M = N*2;
const LL INF = 1e18;
struct Node
{
int l,r;
LL sum,tag;//可能加负数
}tr[N<<2];
int h[N],e[M],ne[M],w[N],idx;
int sz[N],son[N],fa[N],dep[N];
int top[N],id[N],nw[N],ts;
int n,q,root;
void add(int a,int b)
{
// e[idx] 为节点原来标号
// ne[idx] 为同一父亲的儿子前驱标号
e[++idx] = b,ne[idx] = h[a],h[a] = idx;
// 一条边保存两次
}
// 求深度,重链
void dfs1(int u,int pa,int depth)
{
sz[u] = 1,dep[u] = depth;// 根节点1的深度为1
// h[u] 是u的最后一个儿子, ne[i]为u的前面一个儿子,0结束(h[u]的初始值为0)
for(int i=h[u];i;i=ne[i])
{
int j = e[i]; // 节点原来标号
if(j==pa) continue;
fa[j] = u;
dfs1(j,u,depth+1);
sz[u] += sz[j];// 子树大小
if(sz[j]>sz[son[u]]) son[u] = j;// son[u]为u的重链儿子
}
}
// tp为当前子树的根节点,重链时则为整个重链的根节点
// 先重链的dfs顺序保存顶点值w
void dfs2(int u,int tp)
{
top[u] = tp,id[u] = ++ts,nw[ts] = w[u];
if(!son[u]) return ;
dfs2(son[u],tp);
for(int i=h[u];i;i=ne[i])
{
int j = e[i]; // 节点原来标号
if(j==son[u]||j==fa[u]) continue;// 跳过重链,父节点
dfs2(j,j);
}
}
void pushup(int u)
{
tr[u].sum = tr[u<<1].sum + tr[u<<1|1].sum;
}
// 建线段树,放置 先重链的dfs顺序保存顶点值w, 求和
void build(int u,int L,int r)
{
if(L==r)
{
tr[u].l=L;
tr[u].r=r;
tr[u].sum=nw[L], tr[u].tag=0;
return ;
}
tr[u].l=L;
tr[u].r=r;
int mid = (L + r) >> 1;
build(u<<1,L,mid),build(u<<1|1,mid+1,r);
pushup(u);
}
void pushdown(int u)
{
// int &x=n 是引用,也就是说传递的是变量n的地址,在函数中如果改变了x的值,
// 那么n所代表的实参的值也同样改变。
auto &root = tr[u],&left = tr[u<<1],&right = tr[u<<1|1];
if(root.tag)
{
left.tag += root.tag;
left.sum += root.tag*(left.r - left.l + 1);
right.tag += root.tag;
right.sum += root.tag*(right.r - right.l + 1);
root.tag = 0;
}
}
void modify(int u,int L,int r,int k)
{
if(L<=tr[u].l&&tr[u].r<=r)
{
tr[u].tag += k;
tr[u].sum += 1ll*k*(tr[u].r - tr[u].l + 1);
return ;
}
pushdown(u);
int mid = (tr[u].l + tr[u].r) >> 1;
if(L<=mid) modify(u<<1,L,r,k);
if(r>mid) modify(u<<1|1,L,r,k);
pushup(u);
}
LL query(int u,int l,int r)
{
if(l<=tr[u].l&&tr[u].r<=r) return tr[u].sum;
pushdown(u);
int mid = (tr[u].l + tr[u].r) >> 1;
LL res = 0;
if(l<=mid) res += query(u<<1,l,r);
if(r>mid) res += query(u<<1|1,l,r);
return res;
}
int lca(int x,int y)
{
cout<<"root="<<root<<", x="<<x<<", y="<<y<<", ";
while(top[x]!=top[y])
{
if(dep[top[x]]<dep[top[y]]) swap(x,y);
// top[x] 的深度大
x = fa[top[x]];
}
cout<<(dep[x]<dep[y]?x:y)<<endl;
return dep[x]<dep[y]?x:y;
}
// x=lca(root,x); 条件下,返回x的儿子
int find(int x)
{
int u = root;
while(top[u]!=top[x])
{
if(fa[top[u]]==x) return top[u];
u = fa[top[u]];
}
return son[x];
}
// 包含x,y 的最小子树的根节点
int LCA(int x,int y)
{
if(dep[x]>dep[y]) swap(x,y);
// x 的深度小, 不存在lca(x,y)==y
// 1. y,root均在x的子树中
// 见原题图一,x=3,root=4.y=5
if(lca(x,y)==x && lca(x,root)==x) return lca(y,root);
//2.root在x的子树中
if(lca(x,root)==x)return x;
//3.root在y的子树中
if(lca(y,root)==y)return y;
// x在root的子树中,y与root的公共祖先==y与x的公共祖先,所以包含x,y 的最小子树的根节点是 root
//4. root 在x,y的路径上
//见原题图一, 5与4 公共祖先为4,lca(5,6)==lca(6,4),所以包含5,6 的最小子树的根节点是 4
if((lca(x,root)==root&&lca(x,y)==lca(y,root))||(lca(y,root)==root&&lca(x,y)==lca(x,root))) return root;
//5. root在包含x,y 的最小子树的上面,root不属于包含x,y 的最小子树
//见原题图一, 2与4,2与6公共祖先均为1,所以包含4,6 的最小子树的根节点是 3
if(lca(root,x)==lca(y,root)) return lca(x,y);
//6. 类似于情况4. root 在x,y的路径上一个节点的下面
if(lca(x,y)!=lca(x,root)) return lca(x,root);
if(lca(x,y)!=lca(y,root)) return lca(y,root);
}
void mo2(int x,int p)
{
if(root==x)
{
modify(1,1,n,p);// 全树加p
return ;
}
int res = lca(root,x);
// 例如 原x=5,y=4, lca(x,y)=4, root=6
if(res!=x) modify(1,id[x],id[x]+sz[x]-1,p);
else
{
int u = find(x); // u为x的直接儿子
// 全树加p, u为根节点的子树减去p
modify(1,1,n,p);
modify(1,id[u],id[u]+sz[u]-1,-p);
// res 为x的情况,见原题图三
// x=3, root=6, 原x=2,y=4, lca(x,y)=3
}
}
LL q2(int x)
{
if(x==root) return query(1,1,n);// x为根节点,整棵树的情况
int res = lca(root,x);
if(res!=x) return query(1,id[x],id[x]+sz[x]-1); // 以 x为根节点的子树
int u = find(x); // u为x的直接儿子
return query(1,1,n) - query(1,id[u],id[u]+sz[u]-1); // 全树减去以 u为根节点的子树
// res 为x的情况,lca(root,x)=x, 见原题图五,root=4,x=3,
}
int main()
{
FILE *fi;
freopen_s(&fi,"CF916Ein.txt","r",stdin);
scanf_s("%d%d",&n,&q);// the number of vertices, the number of queries
for(int i=1;i<=n;i++) scanf_s("%d",w+i);// initial values of the vertices.
for(int i=0;i<n-1;i++)
{
int u,v;scanf_s("%d%d",&u,&v);
add(u,v),add(v,u);// 一条边保存两次
}
root = 1;
dfs1(1,-1,1);
dfs2(1,1);
build(1,1,n);
while(q--)
{
int op,x,y,c;
scanf_s("%d%d",&op,&x);
if(op==1) root = x; // 改变根节点
else if(op==2)
{
scanf_s("%d%d",&y,&c);
mo2(LCA(x,y),c);
}
else printf("%lld\n",q2(x)); // 以x为根节点的子树顶点值之和
}
return 0;
}