CF809E Surprise me! 题解

题目描述

给定一棵树和排列 \(a_i\) ，求：

\[\frac1{n(n-1)}\sum_{i=1}^n\sum_{j\neq i}\varphi(a_i\cdot a_j)\cdot dis(i,j) \]

其中 \(dis(i,j)\) 为点 \(i,j\) 的树上距离。

数据范围

• \(2\le n\le 2\cdot 10^5\) 。

时间限制 \(\texttt{8s}\) ，空间限制 \(\texttt{256MB}\) 。

分析

\(i,j\) 纠缠在一起肯定做不了，如果 \(i,j\) 独立就好了。

于是你会发现 \(\varphi(a_i\cdot a_j)\) 放在这里很不顺眼，注意到：

\[\varphi(a\cdot b)=\frac{\varphi(a)\cdot\varphi(b)\cdot\gcd(a,b)}{\varphi(\gcd(a,b))} \]

带回题目要求的式子，忽略掉前面的常数：

\[\sum_{i=1}^n\sum_{j=1}^n\frac{\varphi(a_i)\cdot\varphi(a_j) \cdot\gcd(a_i,a_j)}{\varphi(\gcd(a_i,a_j))}\cdot dis(i,j) \]

枚举 \(\gcd\) ：

\[\sum_{d=1}^n\frac d{\varphi(d)}\sum_{i=1}^n\sum_{j=1}^n [\gcd(a_i,a_j)=d]\varphi(a_i)\cdot\varphi(a_j)\cdot dis(i,j) \]

莫比乌斯反演：

\[\sum_{d=1}^n\frac d{\varphi(d)}\sum_{d|x}\mu(\frac xd) \sum_{i=1}^n\sum_{j=1}^n[x|\gcd(a_i,a_j)]\varphi(a_i)\cdot\varphi(a_j)\cdot dis(i,j) \]

交换求和号：

\[\sum_{x=1}^n\sum_{d|x}\frac d{\varphi(d)}\cdot\mu(\frac xd) \sum_{x|a_i}\sum_{x|a_j}\varphi(a_i)\cdot\varphi(a_j)\cdot dis(i,j) \]

关于 \(d\) 的部分显然可以调和级数预处理，重点是后面关于 \(i,j\) 的部分。

记第 \(i\) 个点的权值 \(v_i=\varphi(a_i)\) ，目标计算：

\[\sum_{x|a_i}\sum_{x|a_j}v_i\cdot v_j\cdot dis(i,j) \]

涉及到的点数只有调和级数 \(\mathcal O(n\log n)\) ，考虑建虚树做树形 \(\texttt{dp}\) 。

很自然地把 \(dis(i,j)\) 拆掉：

\[\sum_{x|a_i}\sum_{x|a_j}v_i\cdot v_j\cdot(dep_i+dep_j-2\cdot dep_{lca(i,j)}) \]

然后稍微化简一下：

\[2\cdot\bigg(\sum_{x|a_i}(v_i\cdot dep_i)\sum_{x|a_j}v_j- \sum_{x|a_i}\sum_{x|a_j}v_i\cdot v_j\cdot dep_{lca(i,j)}\bigg) \]

所有部分都可以树形 \(\texttt{dp}\) 解决。

总点数 \(\mathcal O(n\log n)\) ，建虚树自带一只 \(\log\) ，时间复杂度 \(\mathcal O(n\log^2n)\) 。

#include<bits/stdc++.h>
using namespace std;
const int maxn=2e5+5,mod=1e9+7;
int n,u,v,cnt,now,sum,res,top;
int a[maxn],w[maxn],val[maxn];
int b[maxn],p[maxn],mu[maxn],phi[maxn];
vector<int> d[maxn],g[maxn],vec[maxn];
int dep[maxn],dfn[maxn],fa[maxn][18];
int st[maxn],sz[maxn];
bool vis[maxn];
int qpow(int a,int k)
{
    int ans=1;
    while(k)
    {
        if(k&1) ans=1ll*ans*a%mod;
        a=1ll*a*a%mod,k/=2;
    }
    return ans;
}
void init(int n)
{
    mu[1]=1,phi[1]=1;
    for(int i=2,cnt=0;i<=n;i++)
    {
        if(!b[i]) p[++cnt]=i,mu[i]=-1,phi[i]=i-1;
        for(int j=1;j<=cnt&&i*p[j]<=n;j++)
        {
            b[i*p[j]]=1;
            if(i%p[j]==0)
            {
                phi[i*p[j]]=phi[i]*p[j];
                break;
            }
            mu[i*p[j]]=mu[i]*mu[p[j]],phi[i*p[j]]=phi[i]*phi[p[j]];
        }
    }
    for(int i=1;i<=n;i++)
        for(int j=i;j<=n;j+=i)
        {
            d[j].push_back(i);
            w[j]=(w[j]+1ll*i*qpow(phi[i],mod-2)%mod*mu[j/i])%mod;
        }
}
void dfs1(int u,int father)
{
    dfn[u]=++cnt;
    for(auto v:g[u])
    {
        if(v==father) continue;
        dep[v]=dep[u]+1,fa[v][0]=u;
        for(int j=1;j<=17;j++) fa[v][j]=fa[fa[v][j-1]][j-1];
        dfs1(v,u);
    }
}
bool cmp(int a,int b)
{
    return dfn[a]<dfn[b];
}
int lca(int u,int v)
{
    if(dep[u]<dep[v]) swap(u,v);
    for(int i=17;i>=0;i--)
        if(dep[fa[u][i]]>=dep[v])
            u=fa[u][i];
    if(u==v) return u;
    for(int i=17;i>=0;i--)
        if(fa[u][i]!=fa[v][i])
            u=fa[u][i],v=fa[v][i];
    return fa[u][0];
}
void addedge(int u,int v)
{
    g[u].push_back(v);
}
void dfs2(int u)
{
    if(vis[u]) sz[u]=val[u],now=(now+1ll*val[u]*val[u]*dep[u])%mod;
    else sz[u]=0;
    for(auto v:g[u])
    {
        dfs2(v);
        now=(now+2ll*sz[u]*sz[v]%mod*dep[u])%mod;
        sz[u]=(sz[u]+sz[v])%mod;
    }
}
int main()
{
    scanf("%d",&n),init(n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&a[i]),val[i]=phi[a[i]];
        for(auto x:d[a[i]]) vec[x].push_back(i);
    }
    for(int i=1;i<=n-1;i++)
    {
        scanf("%d%d",&u,&v);
        addedge(u,v),addedge(v,u);
    }
    dep[1]=1,dfs1(1,0);
    for(int x=1;x<=n;x++)
    {
        sort(vec[x].begin(),vec[x].end(),cmp);
        now=sum=0;
        for(auto u:vec[x]) vis[u]=true,sum=(sum+1ll*val[u]*dep[u])%mod;
        g[1].clear(),st[top=1]=1;
        for(auto u:vec[x])
        {
            if(u==1) continue;
            int p=lca(u,st[top]);
            if(p!=st[top])
            {
                while(dfn[p]<dfn[st[top-1]]) addedge(st[top-1],st[top]),top--;
                if(p==st[top-1]) addedge(p,st[top--]);
                else g[p].clear(),addedge(p,st[top--]),st[++top]=p;
            }
            g[u].clear(),st[++top]=u;
        }
        while(top>=2) addedge(st[top-1],st[top]),top--;
        dfs2(1);
        for(auto u:vec[x]) vis[u]=false;
        now=(1ll*sz[1]*sum-now)%mod;
        res=(res+2ll*w[x]*now)%mod;
    }
    printf("%lld\n",1ll*(res+mod)*qpow(n*(n-1ll)%mod,mod-2)%mod);
    return 0;
}