题解:AT_abc392_g [ABC392G] Fine Triplets

· · 题解

[ABC392G] Fine Triplets

前言

赛后才知道有官方库。

正文

题意就是问有多少三元组 (i,j,k) 满足 a_j-a_i=a_k-a_j(a_i<a_j<a_k)

移位后得 a_i+a_k=2a_j,将 a_i 看成 x^{a_i},因为 x^ix^j=x^{i+j},于是就可以直接用快速傅里叶变换(多项式乘法)做。2a_i 次项的系数为 v,那么 \frac{v-1}{2} 就是 B=a_i 时的方案数(因为 a_i\ge a_j\ge a_k 的情况要剪掉),加起来即可。

提示:FFT 卷积在官方库中有。

AC Code

#include <atcoder/convolution>
#include <bits/stdc++.h>
#define int long long
using namespace std;
using namespace atcoder;
int a[1000005];
vector <int> s , c , s2;
signed main ()
{
    ios::sync_with_stdio (0) , cin.tie (0) , cout.tie (0);
    int n;
    cin >> n;
    s.resize (1e6 + 5); 
    for (int i = 1;i <= n;i ++)
        cin >> a[i] , s[a[i]] = 1;
    s2 = s;
    c = convolution (s2 , s2);
    int ans = 0;
    for (int i = 1;i <= n;i ++)
        ans += (c[a[i] << 1] - 1) >> 1;
    cout << ans;
    return 0;
}