题目描述
有n根木棒,现在从中选 4 根,想要组成一个正三角形,问有几种选法?
答案对 998244353取模。
输入格式
第一行一个整数 n。
第二行 n 个整数,第
i
i
i个整数
a
i
a_i
ai 代表第 i 根木棒的长度。
输出格式
一行一个整数代表答案。
输入输出样例
输入
4
1 1 2 2
输出
1
说明/提示
数据规模与约定
0
≤
a
i
≤
5
×
1
0
6
,
1
≤
n
≤
1
0
5
0≤a_i≤5×10^{6}, 1≤n≤10^5
0≤ai≤5×106,1≤n≤105
解题思路
对于构建边长为
x
x
x的正三角形,边长为
x
x
x的木棍有
k
k
k根,两根组合成一根长为
x
x
x 的木棍有
y
y
y 种,那么构建边长为x的正三角形有
y
×
C
k
2
y\times C{}^{2}_{k}
y×Ck2
如何计算
y
y
y呢?
设
a
1
,
a
2
,
.
.
.
,
a
m
a_1, a_2, ..., a_m
a1,a2,...,am表示长为
i
i
i的木棍有多少根
对于两根木棍拼起来长为
x
x
x的种类数分为以下两种情况
当
x
x
x为奇数时
y
=
a
1
×
a
x
−
1
+
a
2
×
a
x
−
2
+
.
.
.
+
a
x
/
2
×
a
x
/
2
+
1
y = a_1\times a_{x-1} + a_2\times a_{x-2} + ... + a_{x/2} \times a_{x/2+1}
y=a1×ax−1+a2×ax−2+...+ax/2×ax/2+1
当
x
x
x为偶数时
y
=
a
1
×
a
x
−
1
+
a
2
×
a
x
−
2
+
.
.
.
+
a
x
/
2
×
(
a
x
/
2
−
1
)
/
2
y = a_1\times a_{x-1} + a_2\times a_{x-2} + ... + a_{x/2} \times (a_{x/2} - 1)/ 2
y=a1×ax−1+a2×ax−2+...+ax/2×(ax/2−1)/2
用以上方法可以在
O
(
n
∗
m
)
O(n*m)
O(n∗m)的时间内算出结果,这是基础版的解法,我们进行了数据的加强,我们需要对y的计算进行一个简化,于是我们想到了多项式乘法,由于多项式乘法可以用NTT优化到
O
(
m
l
o
g
m
)
O(mlogm)
O(mlogm)
我们定义多项式
f
(
x
)
=
a
1
x
+
a
2
x
2
+
.
.
.
+
a
i
x
i
+
.
.
.
+
a
m
x
m
f(x) = a_1x + a_2x^2 + ...+a_ix^i + ... + a_mx^m
f(x)=a1x+a2x2+...+aixi+...+amxm
g
(
x
)
=
f
(
x
)
g(x) = f(x)
g(x)=f(x)
h
(
x
)
=
f
(
x
)
×
g
(
x
)
h(x) = f(x) \times g(x)
h(x)=f(x)×g(x)
我们发现,对
h
(
x
)
h(x)
h(x)中的任意一项
b
i
x
i
b_ix^i
bixi进行研究
其中
b
i
=
a
1
×
a
i
−
1
+
a
2
×
a
i
−
2
+
.
.
.
+
a
i
−
1
×
a
1
b_i = a_1 \times a_{i-1} + a_2\times a_{i-2} + ...+ a_{i-1} \times a_1
bi=a1×ai−1+a2×ai−2+...+ai−1×a1
这样一看,我们就可以将两根木棍拼起来长为
x
x
x 的种类数
y
y
y 与
b
i
b_i
bi 联系起来。
当
x
x
x 为奇数时
y
=
b
i
/
2
y = b_i / 2
y=bi/2
当
x
x
x 为偶数时
y
=
(
b
i
−
a
x
/
2
×
a
x
/
2
)
/
2
+
a
x
/
2
×
(
a
x
/
2
−
1
)
/
2
y = (b_i - a_{x/2} \times a_{x/2}) / 2 + a_{x/2} \times (a_{x/2} - 1)/ 2
y=(bi−ax/2×ax/2)/2+ax/2×(ax/2−1)/2
所以我们只需要以
O
(
m
l
o
g
m
)
O(mlogm)
O(mlogm)的复杂度进行一个预处理就可以求出所有两根木棍拼起来长为
x
x
x 的种类数
y
y
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;
}