2022“杭电杯”中国大学生算法设计超级联赛（7）

1003 Counting Stickmen

思路：组合数学，枚举火柴人的脖子和身体 ，注意去重

#include<iostream>
#include<cstring>
#include<algorithm>
#define int long long
using namespace std;
typedef long long ll;
const int N=5e5+10,M=2*N,mod=998244353;
int cnt[N],ng[N],ngc[N];
int e[M],ne[M],h[N],idx;
int n;
struct Edge
{
    int a,b;
}edge[N];
void add(int a,int b)
{
    e[idx]=b;
    ne[idx]=h[a];
    h[a]=idx++;
}
int res;
int calc(int x)
{
    return x*(x-1)/2%mod;
}
void solve()
{
    idx=res=0;
    cin>>n;
    for(int i=1;i<=n;i++)
        h[i]=-1,cnt[i]=0,ng[i]=0,ngc[i]=0;
    for(int i=1;i<n;i++)
    {
        int a,b;
        scanf("%d %d",&a,&b);
        edge[i]={a,b};
        add(a,b);
        add(b,a);
        cnt[a]++;
        cnt[b]++;
    }
    for(int i=1;i<=n;i++)
    {
        for(int j=h[i];j!=-1;j=ne[j])
        {
            int u=e[j];
            ng[i]+=cnt[u]-1;
            ngc[i]=(ngc[i]+calc(cnt[u]-1))%mod;
        }
        //cout<<i<<" "<<ng[i]<<" "<<ngc[i]<<endl;
    }
    for(int i=1;i<n;i++)
    {
        int x=edge[i].a,y=edge[i].b;
        int foot=calc(cnt[y]-1);
        int hand=calc(ng[x]-cnt[y]+1);
        hand=hand-ngc[x]+calc(cnt[y]-1);
        int head=cnt[x]-3;
       // if(hand<0||head<0||foot<0);
       // else
        res=(res+head*hand%mod*foot%mod)%mod;
      //  cout<<x<<" "<<y<<"hand"<<hand<<"head"<<head<<"foot"<<foot<<endl;

        x=edge[i].b,y=edge[i].a;
        foot=calc(cnt[y]-1);
        hand=calc(ng[x]-cnt[y]+1);
        hand=hand-ngc[x]+calc(cnt[y]-1);
        head=cnt[x]-3;
      //  if(hand<0||head<0||foot<0);
      //  else
        res=(res+head*hand%mod*foot%mod)%mod;
      //  cout<<x<<" "<<y<<"hand"<<hand<<"head"<<head<<"foot"<<foot<<endl;
    }
    cout<<res<<endl;
}
signed main()
{
    int t;
    t=1;
    cin>>t;
    for(int i=1;i<=t;i++)
    {
        solve();
    }
    return 0;
}