Silverwolfnya

CF2241E

solutions

原题链接

首先，考虑三个点构成的路径集有且仅有以下两种情况：

  1. 链
  2. 路径两两不交，且\(LCA(u,v,w)=x\)。

如下图所示：

image

对于链的情况，仅有中间的这个节点被覆盖了奇数次，因此乘积是否为完全平方数仅与 \(v\) 节点有关。不妨枚举 $ v$ ，然后在两颗不同的子树选取这样的 \(v,w\) 。这坨东西的贡献实际是好算的。设当前 \(v\) 是整颗树的根，子节点个数为\(n\) 个。

考虑正着算：\(\sum_{i=1}^n \sum_{j=i+1}^n sz_i\times sz_j\)，这个东西是很好前缀和优化的。

考虑容斥算，那就是从其他节点随便选两个，去除掉在一颗子树的贡献。那就是\(\binom{n-1}{2}-\sum_{i=1}^n \binom{sz_i}{2}\)。

而对于第二种情况，仍然钦定 \(x\) 为根，那么答案就是 \(\sum_{i=1}^n \sum_{j=i+1}^n \sum_{k=j+1}^n sz_i\times sz_j\times sz_k\)，依旧是可以前缀和优化的。

于是我们是有 \(O(n)\) 的做法的。

#include<bits/stdc++.h>
#define pb push_back
#define fi first 
#define se second
#define mt make_tuple
using namespace std;
typedef pair<int,int>pii;
typedef tuple<int,int,int>ti3;
typedef tuple<int,int,int,int>ti4;
typedef tuple<int,int,int,int,int>ti5;
using ll=long long;
const int N=1e6+3;
bool is_sq[N];
void init(){
    for(ll i=1;i*i<=N;i++){
        is_sq[i*i]=true;
    }
}
ll C3(ll x){
    if(x<=2)return 0;
    return 1ll*x*(x-1)*(x-2)/6;
}
ll C2(ll x){
    if(x<=1)return 0;
    return x*(x-1)/2;
}
void Silverwolf(){
    int n;
    cin>>n;
    vector<int>a(n+1);
    vector<vector<int>>e(n+1);
    for(int i=1;i<=n;i++)cin>>a[i];
    for(int i=1;i<n;i++){
        int u,v;
        cin>>u>>v;
        e[u].pb(v);
        e[v].pb(u);
    }
    ll ans=0;
    //第一种情况 链
    vector<int>sz(n+1,0);
    auto dfs=[&](auto &&dfs,int u,int f)->void{
        sz[u]=1;
        vector<int>vt;
        for(int v:e[u]){
            if(v==f)continue;
            dfs(dfs,v,u);
            sz[u]+=sz[v];
            if(is_sq[a[u]])ans-=C2(sz[v]);
            vt.push_back(sz[v]);
        }
        vt.push_back(n-sz[u]);
        if(is_sq[a[u]]){
            ans-=C2(n-sz[u]);
            ans+=C2(n-1);
            if(vt.size()>=3){
                vector<ll>pre(vt.size(),0);//先算前缀和
                vector<ll>pre2(vt.size(),0);//然后算第二坨东西
                for(int i=0;i<vt.size();i++){
                    if(i==0){
                        pre[i]=vt[i];
                    }else pre[i]=pre[i-1]+vt[i];
                    if(i>=1)pre2[i]=pre2[i-1]+vt[i]*pre[i-1];
                    if(i>=2)ans+=1ll*vt[i]*pre2[i-1];
                }
            }
        }

    };
    dfs(dfs,1,0);
    cout<<ans<<'\n';
}
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    init();
    // freopen("Straps.in","r",stdin);
    // freopen("Straps.out","w",stdout);
    int T;cin>>T;while(T--)
    Silverwolf();
    return 0;
}

导航