洛谷 P3799 妖梦拼木棒 加强版

题目描述

有n根木棒，现在从中选 4 根，想要组成一个正三角形，问有几种选法？
答案对 998244353取模。

输入格式

第一行一个整数 n。
第二行 n 个整数，第$i$个整数 $a_i$ 代表第 i 根木棒的长度。

输出格式

一行一个整数代表答案。

输入输出样例

输入
4
1 1 2 2
输出
1
说明/提示
数据规模与约定 $0\le a_i\le 5\times 10^6,1\le n\le 10^5$

解题思路

对于构建边长为$x$的正三角形，边长为$x$的木棍有$k$根，两根组合成一根长为 $x$ 的木棍有 $y$ 种，那么构建边长为x的正三角形有$y\times C{}_k^2$
如何计算$y$呢？
设$a_1,a_2,...,a_m$表示长为$i$的木棍有多少根
对于两根木棍拼起来长为$x$的种类数分为以下两种情况
当$x$为奇数时
$y=a_1\times a_{x-1}+a_2\times a_{x-2}+...+a_{x/2}\times a_{x/2+1}$
当$x$为偶数时
$y=a_1\times a_{x-1}+a_2\times a_{x-2}+...+a_{x/2}\times (a_{x/2}-1)/2$
用以上方法可以在$O(n\ast m)$的时间内算出结果，这是基础版的解法，我们进行了数据的加强，我们需要对y的计算进行一个简化，于是我们想到了多项式乘法，由于多项式乘法可以用NTT优化到$O(mlogm)$
我们定义多项式
$f(x)=a_{1}x+a_2{x}^2+...+a_i{x}^i+...+a_m{x}^m$
$g(x)=f(x)$
$h(x)=f(x)\times g(x)$
我们发现，对$h(x)$中的任意一项$b_i{x}^i$进行研究
其中$b_i=a_1\times a_{i-1}+a_2\times a_{i-2}+...+a_{i-1}\times a_1$
这样一看，我们就可以将两根木棍拼起来长为 $x$ 的种类数 $y$ 与 $b_i$ 联系起来。
当 $x$ 为奇数时
$y=b_i/2$
当 $x$ 为偶数时
$y=(b_i-a_{x/2}\times a_{x/2})/2+a_{x/2}\times (a_{x/2}-1)/2$
所以我们只需要以$O(mlogm)$的复杂度进行一个预处理就可以求出所有两根木棍拼起来长为 $x$ 的种类数 $y$
然后跑一遍就可以出结果了

#include <bits/stdc++.h>
#define ll long long
#define qc ios::sync_with_stdio(false); cin.tie(0);cout.tie(0)
#define fi first
#define se second
#define PII pair<int, int>
#define PLL pair<ll, ll>
#define pb push_back
using namespace std;
const int MAXN = 2e6 + 7;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while (!isdigit(ch)){if (ch=='-') f=-1;ch=getchar();}
    while (isdigit(ch)){x=x*10+ch-48;ch=getchar();}
    return x*f;
}
const int G = 3,mod = 998244353;
int n,m,L,R[MAXN];
int A[MAXN],B[MAXN];
ll b[MAXN];
int qpow(int a,int b){
    int ans = 1;
    while(b){
        if(b&1)
            ans = 1ll * ans * a % mod;
        a = 1ll * a * a % mod;
        b >>= 1;
    }
    return ans;
}
void NTT(int* a,int f){
    for (int i = 0; i < n; i++) if (i < R[i]) swap(a[i],a[R[i]]);
    for (int i = 1; i < n; i <<= 1){
        int gn = qpow(G,(mod - 1) / (i << 1));
        for (int j = 0; j < n; j += (i << 1)){
            int g = 1;
            for (int k = 0; k < i; k++,g = 1ll * g * gn % mod){
                int x = a[j + k],y = 1ll * g * a[j + k + i] % mod;
                a[j + k] = (x + y) % mod; a[j + k + i] = (x - y + mod) % mod;
            }
        }
    }
    if (f == 1) return;
    int nv = qpow(n,mod - 2); reverse(a + 1,a + n);
    for (int i = 0; i < n; i++) a[i] = 1ll * a[i] * nv % mod;
}

void solve(){
    cin >> n;
    int maxx = -1;
    m = n;
    for(int i = 1; i <= n; i++){
        int x;
        cin >> x;
        maxx = max(maxx, x);
        b[x]++;
    }
    for(int i = 0; i <= maxx; i++)
        A[i] = B[i] = b[i];
    m = n + m; for (n = 1; n <= m; n <<= 1) L++;
    for (int i = 0; i < n; i++) R[i] = (R[i >> 1] >> 1) | ((i & 1) << (L - 1));
    NTT(A,1); NTT(B,1);
    for (int i = 0; i < n; i++) A[i] = 1ll * A[i] * B[i] % mod;
    NTT(A,-1);

    for(int i = 1; i < MAXN; i++){
        if(i % 2 == 0){
            A[i] = 1ll * ((A[i] - (b[i/2] * b[i/2] % mod) + mod) % mod * qpow(2, mod - 2) % mod + b[i/2] * ((b[i/2] - 1 + mod) % mod) % mod * qpow(2, mod - 2) % mod) % mod;
        }
        else
            A[i] = 1ll * A[i] * qpow(2, mod - 2) % mod;
    }

    ll ans = 0;
    for(int i = 1; i <= maxx; i++){
        ans += 1ll * b[i] * ((b[i] - 1 + mod) % mod) % mod * qpow(2, mod - 2) % mod * A[i] % mod;
		ans %= mod;
    }
    cout << ans << endl;
}

int main()
{
    #ifdef ONLINE_JUDGE
    #else
       freopen("in.txt", "r", stdin);
       freopen("out.txt", "w", stdout);
    #endif

    qc;
    int T;
    // cin >> T;
    T = 1;
    while(T--){

        solve();
    }
    return 0;
}