CF809E Surprise me!

Posted 2020-11-09 Cyhlnj

题意

给你一棵\(n\) 个点的树，每个点有权值 \(a_i\)?
\(a\) 为一个排列
求
\[\frac{1}{n(n?1)}?\sum_{i=1}^{n}\sum_{j=1}^{n}\varphi(a_ia_j)dist_{i,j}\]
\(n≤200000\)
答案对 \(10^9+7\) 取模

Sol

莫比乌斯反演
然后要求
\[\sum_{d=1}^{n}\frac{d}{\varphi(d)}\sum{i=1}{\lfloor\frac{n}{d}\rfloor}\mu(i)F(i*d)\]
其中
\(F(d)=\sum_{d|a_i}\sum_{d|a_j}\varphi(a_i)\varphi(a_j)dist(i,j)\)

直接把\(F(d)\)求出来，把满足要求的点拿出来建虚树\(DP\)就好了

# include <bits/stdc++.h>
# define IL inline
# define RG register
# define Fill(a, b) memset(a, b, sizeof(a))
using namespace std;
typedef long long ll;

IL int Input(){
    RG int x = 0, z = 1; RG char c = getchar();
    for(; c < '0' || c > '9'; c = getchar()) z = c == '-' ? -1 : 1;
    for(; c >= '0' && c <= '9'; c = getchar()) x = (x << 1) + (x << 3) + (c ^ 48);
    return x * z;
}

const int maxn(2e5 + 5);
const int mod(1e9 + 7);

typedef int Arr[maxn];

int n, cnt, idx, tot, len, ret;
Arr dfn, size, fa, first1, first2, s, deep, id;
int lg[maxn << 1], st[20][maxn << 1];
Arr a, prime, isprime, phi, mu, f, mark, sum;
vector <int> p[maxn], fac[maxn];

struct Edge{
    int to, next;
} e1[maxn << 1], e2[maxn];

IL int Pow(RG ll x, RG ll y){
    RG ll ret = 1;
    for(; y; y >>= 1, x = x * x % mod)
        if(y & 1) ret = ret * x % mod;
    return ret;
}

IL void Add1(RG int u, RG int v){
    e1[cnt] = (Edge){v, first1[u]}, first1[u] = cnt++;
}

IL void Add2(RG int u, RG int v){
    e2[cnt] = (Edge){v, first2[u]}, first2[u] = cnt++;
}

IL void Dfs(RG int u){
    size[u] = 1, dfn[u] = ++idx, st[0][id[u] = ++len] = u;
    for(RG int i = 0, j = fac[a[u]].size(); i < j; ++i)
        p[fac[a[u]][i]].push_back(u);
    for(RG int e = first1[u]; e != -1; e = e1[e].next){
        RG int v = e1[e].to;
        if(!size[v]){
            deep[v] = deep[u] + 1;
            Dfs(v), size[u] += size[v];
            fa[v] = u, st[0][++len] = u;
        }
    }
}

IL int Min(RG int x, RG int y){
    return deep[x] < deep[y] ? x : y;
}

IL int LCA(RG int u, RG int v){
    u = id[u], v = id[v];
    if(u > v) swap(u, v);
    RG int lgn = lg[v - u + 1];
    return Min(st[lgn][u], st[lgn][v - (1 << lgn) + 1]);
}

IL int Dis(RG int u, RG int v){
    RG int lca = LCA(u, v);
    return deep[u] + deep[v] - 2 * deep[lca];
}

IL void Upd(RG int &x, RG int y){
    x += y;
    if(x >= mod) x -= mod;
}

IL void DP(RG int u){
    sum[u] = mark[u] * phi[a[u]];
    Upd(ret, mod - 1LL * sum[u] * sum[u] % mod * deep[u] % mod);
    for(RG int e = first2[u]; e != -1; e = e2[e].next){
        RG int v = e2[e].to;
        DP(v);
        Upd(ret, mod - 2LL * sum[u] * sum[v] % mod * deep[u] % mod);
        Upd(sum[u], sum[v]);
    }
}

IL int Calc(RG int x){
    RG int l = p[x].size(), ret1 = 0, ret2 = 0, t = 0; cnt = 0;
    for(RG int i = 0; i < l; ++i) first2[p[x][i]] = -1, mark[p[x][i]] = 1;
    for(RG int i = 0; i < l; ++i){
        RG int lca = s[t] ? LCA(s[t], p[x][i]) : 0;
        if(lca == s[t]) s[++t] = p[x][i];
        else{
            while(t > 1 && deep[lca] <= deep[s[t - 1]]) Add2(s[t - 1], s[t]), --t;
            if(s[t] != lca) first2[lca] = -1, Add2(lca, s[t]), s[t] = lca;
            s[++t] = p[x][i];
        }
    }
    while(t > 1) Add2(s[t - 1], s[t]), --t;
    for(RG int i = 0; i < l; ++i){
        Upd(ret1, phi[a[p[x][i]]]);
        Upd(ret2, 1LL * deep[p[x][i]] * phi[a[p[x][i]]] % mod);
    }
    ret = 0;
    DP(s[1]), Upd(ret, ret);
    Upd(ret, 2LL * ret1 * ret2 % mod);
    for(RG int i = 0; i < l; ++i) mark[p[x][i]] = 0;
    return ret;
}

IL void Sieve(){
    phi[1] = mu[1] = 1;
    for(RG int i = 2; i <= n; ++i){
        if(!isprime[i]) prime[++tot] = i, phi[i] = i - 1, mu[i] = -1;
        for(RG int j = 1; j <= tot && i * prime[j] <= n; ++j){
            isprime[i * prime[j]] = 1;
            if(i % prime[j]){
                phi[i * prime[j]] = phi[i] * phi[prime[j]];
                mu[i * prime[j]] = -mu[i];
            }
            else{
                phi[i * prime[j]] = phi[i] * prime[j];
                mu[i * prime[j]] = 0;
                break;
            }
        }
    }
    for(RG int i = 1; i <= n; ++i)
        for(RG int j = 1; j * i <= n; ++j) fac[j * i].push_back(i);
}

int main(){
    n = Input(), Sieve();
    for(RG int i = 1; i <= n; ++i) first1[i] = first2[i] = -1, a[i] = Input();
    for(RG int i = 1; i < n; ++i){
        RG int u = Input(), v = Input();
        Add1(u, v), Add1(v, u);
    }
    Dfs(1);
    for(RG int i = 2; i <= len; ++i) lg[i] = lg[i >> 1] + 1;
    for(RG int j = 1; j <= lg[len]; ++j)
        for(RG int i = 1; i + (1 << j) - 1 <= len; ++i)
            st[j][i] = Min(st[j - 1][i], st[j - 1][i + (1 << (j - 1))]);
    for(RG int i = 1; i <= n; ++i) f[i] = Calc(i);
    RG int ans = 0;
    for(RG int i = 1; i <= n; ++i){
        RG int ret = 0;
        for(RG int j = 1, t = n / i; j <= t; ++j)
            Upd(ret, 1LL * (mu[j] + mod) % mod * f[i * j] % mod);
        ret = 1LL * ret * i % mod;
        Upd(ans, 1LL * ret * Pow(phi[i], mod - 2) % mod);
    }
    ans = 1LL * ans * Pow(1LL * n * (n - 1) % mod, mod - 2) % mod;
    printf("%d\n", ans);
    return 0;
}