BZOJ 5379 Tree

题目链接

https://www.lydsy.com/JudgeOnline/problem.php?id=5379

题解

显然，换根操作只需要记录，树实际上的根始终都是$1$。

对于修改操作，$a,b$在树上的lca，就是$lca(a,b)$，$lca(a,root)$，$lca(b,root)$中深度最大的一个。

使用线段树+dfs修改一个子树的权值，以及root在不在一个点的子树内（理论上来说可以用树链剖分）

代码

#include <cstdio>
#include <algorithm>

const int maxn=300000;

int read()
{
  int x=0,f=1;
  char ch=getchar();
  while((ch<'0')||(ch>'9'))
    {
      if(ch=='-')
        {
          f=-f;
        }
      ch=getchar();
    }
  while((ch>='0')&&(ch<='9'))
    {
      x=x*10+ch-'0';
      ch=getchar();
    }
  return x*f;
}

long long adt,w[maxn+10];
int fa[maxn+10][22],dfn[maxn+10],size[maxn+10];
int cnt,n,q,nowroot,deep[maxn+10],idfn[maxn+10];

namespace segment_tree
{
  long long val[maxn<<2],lazy[maxn<<2];

  int putadt(int x,int l,int r,long long v)
  {
    val[x]+=1ll*(r-l+1)*v;
    lazy[x]+=v;
    return 0;
  }

  int pushdown(int now,int l,int r)
  {
    int mid=(l+r)>>1;
    putadt(now<<1,l,mid,lazy[now]);
    putadt(now<<1|1,mid+1,r,lazy[now]);
    lazy[now]=0;
    return 0;
  }

  int updata(int now)
  {
    val[now]=val[now<<1]+val[now<<1|1];
    return 0;
  }

  int build(int now,int l,int r)
  {
    if(l==r)
      {
        val[now]=w[idfn[l]];
        return 0;
      }
    int mid=(l+r)>>1;
    build(now<<1,l,mid);
    build(now<<1|1,mid+1,r);
    updata(now);
    return 0;
  }

  int add(int now,int l,int r,int askl,int askr,int adv)
  {
    if((askl<=l)&&(r<=askr))
      {
        putadt(now,l,r,adv);
        return 0;
      }
    int mid=(l+r)>>1;
    pushdown(now,l,r);
    if(askl<=mid)
      {
        add(now<<1,l,mid,askl,askr,adv);
      }
    if(mid<askr)
      {
        add(now<<1|1,mid+1,r,askl,askr,adv);
      }
    updata(now);
    return 0;
  }

  long long getsum(int now,int l,int r,int askl,int askr)
  {
    if((askl<=l)&&(r<=askr))
      {
        return val[now];
      }
    int mid=(l+r)>>1;
    long long res=0;
    pushdown(now,l,r);
    if(askl<=mid)
      {
        res+=getsum(now<<1,l,mid,askl,askr);
      }
    if(mid<askr)
      {
        res+=getsum(now<<1|1,mid+1,r,askl,askr);
      }
    return res;
  }
}

int pre[maxn*2+10],now[maxn+10],son[maxn*2+10],tot;

int ins(int a,int b)
{
  pre[++tot]=now[a];
  now[a]=tot;
  son[tot]=b;
  return 0;
}

int dfs(int u,int f)
{
  fa[u][0]=f;
  dfn[u]=++cnt;
  size[u]=1;
  idfn[cnt]=u;
  deep[u]=deep[f]+1;
  int j=now[u];
  while(j)
    {
      int v=son[j];
      if(v!=f)
        {
          dfs(v,u);
          size[u]+=size[v];
        }
      j=pre[j];
    }
  return 0;
}

int getfa()
{
  for(int j=1; j<=20; ++j)
    {
      for(int i=1; i<=n; ++i)
        {
          fa[i][j]=fa[fa[i][j-1]][j-1];
        }
    }
  return 0;
}

int getlca(int a,int b)
{
  if(deep[a]<deep[b])
    {
      std::swap(a,b);
    }
  for(int i=20; ~i; --i)
    {
      if(deep[fa[a][i]]>=deep[b])
        {
          a=fa[a][i];
        }
    }
  if(a==b)
    {
      return a;
    }
  for(int i=20; ~i; --i)
    {
      if(fa[a][i]!=fa[b][i])
        {
          a=fa[a][i];
          b=fa[b][i];
        }
    }
  return fa[a][0];
}

int ris(int u)
{
  return (dfn[u]<=dfn[nowroot])&&(dfn[nowroot]<dfn[u]+size[u]);
}

int gts(int u,int f)
{
  for(int i=20; ~i; --i)
    {
      if(deep[fa[u][i]]>deep[f])
        {
          u=fa[u][i];
        }
    }
  return u;
}

int pushadt(int a,int b,int val)
{
  int lca=getlca(a,b),lcaa=getlca(a,nowroot),lcab=getlca(nowroot,b);
  if(deep[lcaa]>deep[lca])
    {
      lca=lcaa;
    }
  if(deep[lcab]>deep[lca])
    {
      lca=lcab;
    }
  if(ris(lca))
    {
      segment_tree::add(1,1,n,1,n,val);
      if(lca!=nowroot)
        {
          int f=gts(nowroot,lca);
          segment_tree::add(1,1,n,dfn[f],dfn[f]+size[f]-1,-val);
        }
    }
  else
    {
      segment_tree::add(1,1,n,dfn[lca],dfn[lca]+size[lca]-1,val);
    }
  return 0;
}

long long getsum(int u)
{
  long long res;
  if(ris(u))
    {
      res=segment_tree::getsum(1,1,n,1,n);
      if(nowroot!=u)
        {
          int f=gts(nowroot,u);
          res-=segment_tree::getsum(1,1,n,dfn[f],dfn[f]+size[f]-1);
        }
    }
  else
    {
      res=segment_tree::getsum(1,1,n,dfn[u],dfn[u]+size[u]-1);
    }
  return res;
}

int main()
{
  n=read();
  q=read();
  for(int i=1; i<=n; ++i)
    {
      w[i]=read();
    }
  for(int i=1; i<n; ++i)
    {
      int a=read(),b=read();
      ins(a,b);
      ins(b,a);
    }
  dfs(1,0);
  getfa();
  nowroot=1;
  segment_tree::build(1,1,n);
  while(q--)
    {
      int op=read(),a=read();
      if(op==1)
        {
          nowroot=a;
        }
      else if(op==2)
        {
          int b=read(),c=read();
          pushadt(a,b,c);
        }
      else
        {
          printf("%lld\n",getsum(a));
        }
    }
  return 0;
}