小清新数据结构题

真是小清新

题目链接

题意

查询以$x$为根时的所有子树的权值和的平方和,带单点修改、

做法

为了表达方便,记$q$为询问的点,我们用$x$表示$1$到$q$路径上的点

可以先$O(n)$的处理出以$1$为根时:

​ 答案,记为$Ans$

​ 每颗子树的权值和,记为$S_i$

​ 所有点权值和,记为$Sum$

​ 点的深度,记为$dep_i$,$dep_1=1$

考虑修改

修改点$q$,记原来值和当前值差为$d$

$Sum=Sum+dep_x*d$

$Ans=Ans-\sum\limits_x(S_x)^2+\sum\limits_x(S_x+d)^2$

$=Ans-\sum\limits_x(S_x)^2+\sum\limits_x(S_x)^2+2S_xd+d^2$

$=Ans+\sum\limits_{x}2S_xd+d^2$

$=Ans+2 d\sum\limits_x S_x+dep_x d^2$

所有$S_x=S_x+d$

当根换为$q$时

此时点$x$,不含点$q$,$S_x$会变为$Sum-S_{son}$,$S_{son}$表示$x$的包含点$q$的子树的大小

其余点不变

所以答案为

$Ans-\sum\limits_x(S_x)^2+Sum^2+\sum\limits_x(Sum-S_x)^2$

$=Ans+\sum\limits_x(Sum-S_x)^2-(S_x)^2+Sum^2$

平方差公式展开

$=Ans+Sum^2+\sum\limits Sum(Sum-2S_x)$

$=Ans+Sum^2+Sum\sum\limits Sum-2S_x$

$=Ans+Sum^2+Sum^2dep_x+2Sum\sum\limits S_x$

用树剖+树状数组维护链上加链上求和就好了

代码

时间复杂度$O(nlog^2n)$

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
/*
>Author: zxy_hhhh
>blog: zxy-hhhh.cn
>date: 2019/01/04
*/
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cctype>
#include<cmath>
#define rep(x,a,b) for (int x=int(a);x<=(int)(b);x++)
#define drp(x,a,b) for (int x=int(a);x>=(int)(b);x--)
#define cross(x,a) for (int x=hd[a];x;x=nx[x])
#define ll long long
using namespace std;
inline ll rd()
{
ll _x=0;int _ch=getchar(),_f=1;
for(;!isdigit(_ch)&&(_ch!='-')&&(_ch!=EOF);_ch=getchar());
if (_ch=='-'){_f=0;_ch=getchar();}
for(;isdigit(_ch);_ch=getchar()) _x=_x*10+_ch-'0';
return _f?_x:-_x;
}
void write(ll _x){if (_x>=10) write(_x/10),putchar(_x%10+'0'); else putchar(_x+'0'); }
inline void wrt(ll _x,char _p){if (_x<0) putchar('-'),_x=-_x; write(_x); if (_p) putchar(_p);}
#define maxn 200005
namespace Tree_Array
{
ll sum[maxn],Sum[maxn];
int n;
inline void add(int x,int val)
{
for(int i=x;i<=n;i+=i&(-i))
sum[i]+=val,Sum[i]+=1ll*x*val;
}
inline ll getsum(int x)
{
ll ans=0;
for(int i=x;i;i-=i&(-i))
ans+=(x+1)*sum[i]-Sum[i];
return ans;
}
inline void update(int l,int r,int x) {add(l,x),add(r+1,-x);}
inline ll query(int l,int r){return getsum(r)-getsum(l-1);}
}
namespace Tree
{
int n;
int nx[maxn<<1],to[maxn<<1],hd[maxn],cnt;
inline void add(int u,int v){nx[++cnt]=hd[u],to[cnt]=v,hd[u]=cnt;}
int top[maxn],sz[maxn],fa[maxn],son[maxn],dep[maxn];
int val[maxn],a[maxn];
int sum[maxn];
ll Ans,S;
int idx[maxn],id;
void dfs(int u)
{
dep[u]=dep[fa[u]]+1,sz[u]=1;sum[u]=a[u];
cross(i,u) if (to[i]!=fa[u]){
fa[to[i]]=u,dfs(to[i]);
sz[u]+=sz[to[i]],sum[u]+=sum[to[i]];
if (sz[son[u]]<sz[to[i]]) son[u]=to[i];
}
}
void dfs(int u,int tp)
{
idx[u]=++id;val[id]=a[u];
top[u]=tp;
if (son[u]) dfs(son[u],tp);
cross(i,u) if (to[i]!=fa[u]&&to[i]!=son[u]) dfs(to[i],to[i]);
}
inline void init()
{
dfs(1),dfs(1,1);
Tree_Array::n=id;
rep(i,1,id) Tree_Array::update(idx[i],idx[i],sum[i]),Ans+=sum[i]*sum[i];
S=sum[1];
}
inline void link_update(int x,ll val)
{
while(1){
Tree_Array::update(idx[top[x]],idx[x],val);
if (top[x]==1) return ;
x=fa[top[x]];
}
}
inline ll sigma(int x)
{
ll ans=0;
while(1){
ans+=Tree_Array::query(idx[top[x]],idx[x]);
if (top[x]==1) return ans;
x=fa[top[x]];
}
}
inline ll query(int x) {return Ans+S*S+1ll*dep[x]*S*S-2ll*S*sigma(x);}
inline void update(int x,ll val)
{
ll delta=val-a[x],s1=sigma(x),s2=dep[x];
Ans+=2*s1*delta+delta*delta*dep[x];
S+=delta;
link_update(x,delta);
a[x]=val;
}
}
using namespace Tree;
int main()
{
n=rd();
int m=rd();
rep(i,1,n-1) {
int x=rd(),y=rd();
add(x,y);add(y,x);
}
rep(i,1,n) a[i]=rd();
init();
rep(_i,1,m){
int op=rd();
if (op==1){
int x=rd(),val=rd();
update(x,val);
}
else {
int x=rd();
wrt(query(x),'\n');
}
}
}
avatar
圆方树

  1. 1. 题意
  2. 2. 做法
    1. 2.1. 可以先$O(n)$的处理出以$1$为根时:
    2. 2.2. 考虑修改
    3. 2.3. 当根换为$q$时
  3. 3. 代码