松鼠的新家 (树链剖分)

为了熟悉树链剖分，又从洛谷上找了一道比较水的题松鼠的新家练手。

题目大意

一只熊要按规定路线在树上走，沿途每走到一个点，都要吃一块糖，问每个点最少需要事先准备多少个糖

思路

正解好像是树剖+树上差分，不过我发现树剖+线段树也能做，不过要记得扣除重复的计数以及最后到的点不计数。总之是树剖裸题

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<map>
#include<queue>
#include<string>
#include<vector>
#include<cmath>
#include<climits>
#include<functional>
#include<set>
#define dd(x) cout<<#x<<" = "<<x<<" "
#define de(x) cout<<#x<<" = "<<x<<endl
#define fi first
#define se second
#define mp make_pair
#define pb push_back
using namespace std;
typedef long long ll;
typedef pair<int,int> P;
typedef vector<int> V;
typedef map<int,int> M;
typedef queue<int> Q;
typedef priority_queue<int> BQ;
typedef priority_queue<int,vector<int>,greater<int> > SQ;
const int maxn=3e5+10,INF=0x3f3f3f3f;
int cnt,head[maxn];
struct Edge
{
  int v,ne;
}edge[maxn<<1];
void add(int u,int v)
{
  edge[++cnt].ne=head[u];
  edge[cnt].v=v;
  head[u]=cnt;
}
int tot,fa[maxn],son[maxn],id[maxn],rk[maxn],top[maxn],sz[maxn],dep[maxn];
void dfs1(int f,int u,int deep)
{
  fa[u]=f;
  dep[u]=deep+1;
  sz[u]=1;
  for (int i=head[u];i;i=edge[i].ne)
  {
    int v=edge[i].v;
    if (v==f)
      continue;
    dfs1(u,v,deep+1);
    sz[u]+=sz[v];
    if (sz[v]>sz[son[u]])
      son[u]=v;
  }
}
void dfs2(int u,int t)
{
  top[u]=t;
  id[u]=++tot;
  rk[tot]=u;
  if (!son[u])
    return;
  dfs2(son[u],t);
  for (int i=head[u];i;i=edge[i].ne)
  {
    int v=edge[i].v;
    if (v!=fa[u]&&v!=son[u])
      dfs2(v,v);
  }
}
ll sum[maxn<<2],lazy[maxn<<2];
void upd(int L,int R,ll c,int l,int r,int rt)
{
  sum[rt]+=c*(R-L+1);
  if (L==l&&R==r)
  {
    lazy[rt]+=c;
    return;
  }
  int m=(l+r)>>1;
  if (R<=m)
    upd(L,R,c,l,m,rt<<1);
  else if (L>m)
    upd(L,R,c,m+1,r,rt<<1|1);
  else 
    upd(L,m,c,l,m,rt<<1),upd(m+1,R,c,m+1,r,rt<<1|1);
}
ll qry(int p,ll add,int l,int r,int rt)
{
  if (l==r)
    return sum[rt]+add;
  int m=(l+r)>>1;
  if (p<=m)
    return qry(p,add+lazy[rt],l,m,rt<<1);
  else
    return qry(p,add+lazy[rt],m+1,r,rt<<1|1);
}
void upd_path(int x,int y)
{
  int fx=top[x],fy=top[y];
  while (fx!=fy)
  {
    if (dep[fx]>dep[fy])
    {
      upd(id[fx],id[x],1,1,tot,1);
      x=fa[fx];
    }
    else
    {
      upd(id[fy],id[y],1,1,tot,1);
      y=fa[fy];
    }
    fx=top[x],fy=top[y];
  }
  if (id[x]<id[y])
    upd(id[x],id[y],1,1,tot,1);
  else
  	upd(id[y],id[x],1,1,tot,1);
}
int a[maxn],path[maxn];
int main()
{
  	int n;
  	scanf("%d",&n);
  	for (int i=0;i<n;++i)
  	{
  		scanf("%d",&path[i]);
  		if (i>0)
  			a[path[i]]++;
  	}
  	for (int i=0;i<n-1;++i)
  	{
  		int u,v;
  		scanf("%d%d",&u,&v);
  		add(u,v);
  		add(v,u);
  	}
  	dfs1(0,1,1);
  	dfs2(1,1);
  	for (int i=1;i<n;++i)
  	{
  		int u=path[i-1],v=path[i];
  		upd_path(u,v);
  	}
  	for (int i=1;i<=n;++i)
  		printf("%lld\n",qry(id[i],0,1,tot,1)-a[i]);
  return 0;
}