P4689 [Ynoi2016] 这是我自己的发明 题解

题目描述

给定一棵 \(n\) 个点的树，点有点权 \(w_i\) ，初始根节点为 \(1\) 号点。

接下来 \(m\) 次操作：

• 1 x ：将根节点换为 \(x\) 号点。
• 2 x y ：询问从 \(x\) 子树和 \(y\) 子树中分别选一个点，点权相等的方案数。

数据范围

• \(1\le n\le 10^5,1\le m\le 5\cdot 10^5,1\le w_i\le 10^9\) 。
• \(1\le x,y\le n\) 。

时间限制 \(\texttt{1.5s}\) ，空间限制 \(\texttt{512MB}\) 。

分析

容易发现换根是假的，钦定 \(1\) 号点为根，那么询问范围只有三种：

• 根节点恰好为 \(x\) ，询问范围为所有点。
• 根节点在 \(x\) 子树外，询问范围为 \(x\) 子树。
• 根节点在 \(x\) 子树内，询问范围为 \([1,n]\) 扣掉 \(x\) 在 \(rt\) 方向上的子树。

用 \(dfs\) 序刻画子树，则询问范围为至多 \(2\) 个区间。

再来考虑如何处理询问，定义 \(f(a,b,c,d)=\sum_{i=a}^b\sum_{j=c}^d[w_i=w_j]\) 。

定义 \(g(a,b)=\sum_{i=1}^a\sum_{j=1}^b[w_i=w_j]\) ，容易发现：

\[f(a,b,c,d)=g(b,d)-g(a-1,d)-g(b,c-1)+g(a-1,c-1) \]

\(g(a,b)\) 显然可以用莫队计算，但直接拆会带上 \(2\times 2\times 4=16\) 倍常数，过不去。

---

继续挖掘题目性质：

• \(f(a,b,c,d)=f(c,d,a,b)\)。
• 如果要拆两个区间，那么第一个区间以 \(1\) 开头，第二个区间以 \(n\) 结尾。

\(a=1\) 是不需要拆的，对于 \(d=n\) ，预处理 \(h(b)=g(n,b)=\sum_{i=1}^n\sum_{j=1}^b[w_i=w_j]\) 。

现在拆区间就容易多了：

• \(f(x_1,y_1,1,x_2)=g(y_1,x_2)-g(x_1-1,x_2)\)。
• \(f(x_1,y_1,y_2,n)=h(y_1)-h(x_1-1)-g(y_1,y_2-1)+g(x_1-1,y_2-1)\)。

这样每个询问会被拆成至多 \(4\) 个关于 \(g\) 的询问，时间复杂度 \(\mathcal O(n\sqrt m+m\log m)\) 。

#include<bits/stdc++.h>
#define ll long long
#define fi first
#define se second
#define mp make_pair
#define pii pair<int,int>
using namespace std;
const int maxn=1e5+5,B=70;
int m,n,q,rt=1,num;
int a[maxn],c[maxn],w[maxn],bel[maxn],cnt[2][maxn];
int d[maxn],sz[maxn],dfn[maxn],fa[maxn][17];
ll cur,h[maxn],res[5*maxn];
vector<int> g[maxn];
struct quer
{
    int l,r,id,sgn;
}f[20*maxn];
inline int read()
{
    int q=0;char ch=getchar();
    while(!isdigit(ch)) ch=getchar();
    while(isdigit(ch)) q=10*q+ch-'0',ch=getchar();
    return q;
}
inline bool cmp(const quer &a,const quer &b)
{
    if(bel[a.l]!=bel[b.l]) return bel[a.l]<bel[b.l];
    return bel[a.l]&1?a.r<b.r:a.r>b.r;
}
void dfs(int u,int f)
{
    dfn[u]=++num,sz[u]=1;
    for(auto v:g[u])
    {
        if(v==f) continue;
        d[v]=d[u]+1,fa[v][0]=u;
        for(int i=1;i<=16;i++) fa[v][i]=fa[fa[v][i-1]][i-1];
        dfs(v,u),sz[u]+=sz[v];
    }
}
inline vector<pii> get(int x)
{
    if(rt==x) return {mp(1,n)};
    if(dfn[rt]<dfn[x]||dfn[rt]>=dfn[x]+sz[x]) return {mp(dfn[x],dfn[x]+sz[x]-1)};
    int u=rt;
    for(int i=16;i>=0;i--) if(d[fa[u][i]]>d[x]) u=fa[u][i];
    return {mp(1,dfn[u]-1),mp(dfn[u]+sz[u],n)};
}
inline void push(pii a,pii b,int id)
{
    if(a.fi>a.se||b.fi>b.se) return ;
    if(a.fi==1) swap(a,b);
    if(b.fi==1)
    {
        f[++q]={a.se,b.se,id,1},f[++q]={a.fi-1,b.se,id,-1};
        return ;
    }
    if(a.se==n) swap(a,b);
    if(b.se==n)
    {
        res[id]+=h[a.se]-h[a.fi-1];
        f[++q]={a.se,b.fi-1,id,-1},f[++q]={a.fi-1,b.fi-1,id,1};
        return ;
    }
    f[++q]={a.se,b.se,id,1},f[++q]={a.fi-1,b.se,id,-1};
    f[++q]={a.se,b.fi-1,id,-1},f[++q]={a.fi-1,b.fi-1,id,1};
}
inline void add(int x,int k)
{
    cur+=cnt[k^1][x],cnt[k][x]++;
}
inline void del(int x,int k)
{
    cur-=cnt[k^1][x],cnt[k][x]--;
}
int main()
{
    n=read(),m=read();
    for(int i=1;i<=n;i++) a[i]=c[i]=read(),bel[i]=(i-1)/B+1;
    for(int i=1;i<=n-1;i++)
    {
        int u=read(),v=read();
        g[u].push_back(v),g[v].push_back(u);
    }
    d[1]=1,dfs(1,0);
    sort(c+1,c+n+1);
    int k=unique(c+1,c+n+1)-c-1;
    for(int i=1;i<=n;i++) cnt[0][w[dfn[i]]=lower_bound(c+1,c+k+1,a[i])-c]++;
    for(int i=1;i<=n;i++) h[i]=h[i-1]+cnt[0][w[i]];
    memset(cnt[0],num=0,sizeof(cnt[0]));
    for(int i=1;i<=m;i++)
    {
        int op=read();
        if(op==1) rt=read();
        else
        {
            num++;
            vector<pii> v1=get(read()),v2=get(read());
            for(auto a:v1) for(auto b:v2) push(a,b,num);
        }
    }
    sort(f+1,f+q+1,cmp);
    for(int i=1,l=0,r=0;i<=q;i++)
    {
        while(l<f[i].l) add(w[++l],0);
        while(l>f[i].l) del(w[l--],0);
        while(r<f[i].r) add(w[++r],1);
        while(r>f[i].r) del(w[r--],1);
        res[f[i].id]+=cur*f[i].sgn;
    }
    for(int i=1;i<=num;i++) printf("%lld\n",res[i]);
    return 0;
}