Codeforces Round #415 (Div. 1), problem: (E) Surprise me! Solution in C/C++

#include<iostream>
#include<vector>
#include<algorithm>
#include<cstring>
#include<cstdio>
#include<cmath>
#include<cstdlib>
#include<ctime>
#include<queue>
#include<set>
using namespace std;
typedef long long LL;
const int N=4e5;
int gi() {
int w=0;bool q=1;char c=getchar();
while ((c<‘0’||c>’9′) && c!=’-‘) c=getchar();
if (c==’-‘) q=0,c=getchar();
while (c>=’0’&&c <= ‘9’) w=w*10+c-‘0’,c=getchar();
return q? w:-w;
}
int head[N],next[N],to[N],fa[N],dfn[N],tot;
int p[N],f[N];
int pri[N],phi[N],G[N],inv[N];
int top[N],dep[N],son[N],siz[N];
int pre[N],st[N],q[N],v[N],sum[N];
inline void dfs1(int k) {
dfn[k]=++tot;siz[k]=1;dep[k]=dep[fa[k]]+1;
for (int i=head[k];i;i=next[i])
if (to[i]!=fa[k]) {
fa[to[i]]=k,dfs1(to[i]),siz[k]+=siz[to[i]];
if (siz[to[i]]>siz[son[k]]) son[k]=to[i];
}
}
inline void dfs2(int k) {
if (son[k]) {
top[son[k]]=top[k];
dfs2(son[k]);
for (int i=head[k];i;i=next[i])
if (to[i]!=fa[k]&&to[i]!=son[k])
top[to[i]]=to[i],dfs2(to[i]);
}
}
inline int lca(int x,int y) {
while (top[x]!=top[y])
dep[top[x]]>dep[top[y]]?x=fa[top[x]]:y=fa[top[y]];
return dep[x]<dep[y]?x:y;
}
inline bool cmp(const int &a,const int &b) { return dfn[a]>dfn[b]; }
int main()
{
#ifndef ONLINE_JUDGE
freopen(“tree.in”,”r”,stdin);
freopen(“tree.out”,”w”,stdout);
#endif
int n=gi(),i,a,b,c,k,tot=0,len=0,t,g,top,j,all;
const int mod=1e9+7;
for (inv[1]=1,i=2;i<=n;i++) inv[i]=mod-1LL*(mod/i)*inv[mod%i]%mod;
for (i=2,phi[1]=G[1]=1;i<=n;i++) {
if (!pri[i]) pri[++len]=i,phi[i]=i-1,G[i]=1LL*i*inv[i-1]%mod;
for (j=1;(t=i*pri[j])<=n;j++) {
pri[t]=1;
if (i%pri[j]) {
phi[t]=phi[i]*phi[pri[j]];
G[t]=1LL*G[i]*G[pri[j]]%mod;
} else {
phi[t]=phi[i]*pri[j];
G[t]=G[i];//!
break;
}
}
}
for (i=1;i<=n;i++) p[k=gi()]=i,v[i]=phi[k];
for (i=1;i<n;i++) {
a=gi(),b=gi();
to[++tot]=b,next[tot]=head[a],head[a]=tot;
to[++tot]=a,next[tot]=head[b],head[b]=tot;
}
dfs1(1);::top[1]=1;dfs2(1);
for (g=1;g<=n;g++) {
for (len=i=0;(i+=g)<=n;)
q[++len]=p[i],sum[p[i]]=v[p[i]];
sort(q+1,q+1+len,cmp);
for (st[top=1]=q[len],i=len;–i;) {
k=q[i];c=lca(k,st[top]);
for (t=0;dep[st[top]]>dep[c];top–)
pre[t=st[top]]=st[top-1];
if (st[top]!=c) {
st[++top]=q[++len]=c;
if (t) pre[t]=c;
}
st[++top]=k;
}
while (–top) pre[st[top+1]]=st[top];
sort(q+1,q+1+len,cmp);//再排序
for (i=1;i<len;i++) {
k=q[i];
f[g]=(f[g]-2LL*sum[pre[k]]*sum[k]%mod*dep[pre[k]])%mod;
(sum[pre[k]]+=sum[k])%=mod;
}
all=sum[q[len]];
for (i=0;(i+=g)<=n;)//只能选倍数
f[g]=(f[g]+1LL*v[p[i]]*dep[p[i]]%mod*(all-v[p[i]]))%mod;
st[++top]=k;
for (i=len;i;i–) sum[q[i]]=0;
}
for (i=n>>1;i;i–)//从后往前
for (j=i<<1;j<=n;j+=i)
(f[i]-=f[j])%=mod;
tot=0;
for (i=1;i<=n;i++)
tot=(tot+1LL*f[i]*G[i])%mod;
tot=1LL*tot*inv[n]%mod*inv[n-1]%mod;
printf(“%d\n”,(((tot+mod)%mod)<<1)%mod);
return 0;
}