题目概述
给出一棵
n
个点的树,每个点有点权
解题报告
naive想法:一步一步跳到LCA,同时统计。
优化:用ST算法倍增来跳,快速统计。
统计的时候记录一下子树和(包括根节点)以及子树平方和(不包括根节点)就行了。
示例程序
#include<cstdio>
#include<algorithm>
using namespace std;
typedef long long LL;
const int maxn=200000,Log=18;
int n,te,a[maxn+5],fa[maxn+5][Log+5],dep[maxn+5];
int E,lnk[maxn+5],son[maxn*2+5],nxt[maxn*2+5];
LL sum[maxn+5],val[maxn+5],Tot[maxn+5][Log+5],ans;
#define Eoln(x) ((x)==10||(x)==13||(x)==EOF)
inline char readc()
{
static char buf[100000],*l=buf,*r=buf;
if (l==r) r=(l=buf)+fread(buf,1,100000,stdin);
if (l==r) return EOF;return *l++;
}
inline int readi(int &x)
{
int tot=0,f=1;char ch=readc(),lst='+';
while ('9'<ch||ch<'0') {if (ch==EOF) return EOF;lst=ch;ch=readc();}
if (lst=='-') f=-f;
while ('0'<=ch&&ch<='9') tot=(tot<<3)+(tot<<1)+ch-48,ch=readc();
return x=tot*f,Eoln(ch);
}
#define sqr(x) ((x)*(x))
void Dfs(int x)
{
sum[x]=a[x];
for (int j=lnk[x];j;j=nxt[j]) if (son[j]!=fa[x][0])
{
fa[son[j]][0]=x;dep[son[j]]=dep[x]+1;Dfs(son[j]);
sum[x]+=sum[son[j]];val[x]+=sqr(sum[son[j]]);
}
for (int j=lnk[x];j;j=nxt[j]) if (son[j]!=fa[x][0])
Tot[son[j]][0]=val[x]-sqr(sum[son[j]]);
}
void Make()
{
for (int j=1;j<=Log;j++)
for (int i=1;i<=n;i++)
fa[i][j]=fa[fa[i][j-1]][j-1],Tot[i][j]=Tot[i][j-1]+Tot[fa[i][j-1]][j-1];
}
#define Add(x,y) son[++E]=(y),nxt[E]=lnk[(x)],lnk[(x)]=E
int main()
{
freopen("mine.in","r",stdin);
freopen("mine.out","w",stdout);
readi(n);readi(te);for (int i=1;i<=n;i++) readi(a[i]);
for (int i=1,x,y;i<n;i++) readi(x),readi(y),Add(x,y),Add(y,x);
for (Dfs(1),Make();te;te--)
{
int x,y;readi(x);readi(y);
if (dep[x]<dep[y]) swap(x,y);ans=val[x];
for (int j=Log;j>=0;j--) if (dep[x]-dep[y]>=(1<<j)) ans+=Tot[x][j],x=fa[x][j];
if (x==y) {printf("%lld\n",ans+sqr(sum[1]-sum[x]));continue;}ans+=val[y];
for (int j=Log;j>=0;j--) if (fa[x][j]!=fa[y][j])
ans+=Tot[x][j]+Tot[y][j],x=fa[x][j],y=fa[y][j];
int lca=fa[x][0];ans+=val[lca]-sqr(sum[x])-sqr(sum[y])+sqr(sum[1]-sum[lca]);
printf("%lld\n",ans);
}
return 0;
}