AtCoder Beginner Contest 392 G

https://atcoder.jp/contests/abc392/tasks/abc392_g

A + C = 2B
なので、A+Cを計算して2Bと合うか調べればよいです。
A+Cは入力例1なら、を計算すればよいです。これはふつうに計算すれば多項式の長さの2乗の計算量ですが、畳み込みはフーリエ変換すれば単なる項目同士の掛け算になります。それを逆フーリエ変換すればよいです。
これで制限時間の約半分の計算時間になりました。

// Fine Triplets
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// FFT ////////////////////

extern crate num_complex;
use num_complex::Complex;

fn fft(x: &Vec<Complex<f64>>) -> Vec<Complex<f64>> {
    let N = x.len();
    if N <= 1 {
        return x.to_vec()
    }
    
    let even = fft(&x.iter().step_by(2).cloned().collect::<Vec<_>>());
    let odd = fft(&x.iter().skip(1).step_by(2).cloned().collect::<Vec<_>>());
    let pi = std::f64::consts::PI;
    let n: f64 = N as f64;
    let T = (0..N/2).map(|k| Complex::new(0.0, -2.0*pi*(k as f64)/n).exp()
                                                * odd[k]).collect::<Vec<_>>();
    let mut v: Vec<Complex<f64>> = vec![Complex::new(0.0, 0.0); N];
    for k in 0..N {
        if k < N / 2 {
            v[k] = even[k] + T[k]
        }
        else {
            v[k] = even[k-N/2] - T[k-N/2]
        }
    }
    v
}

fn ifft(x: &Vec<Complex<f64>>) -> Vec<Complex<f64>> {
    let N = x.len();
    if N <= 1 {
        return x.to_vec()
    }
    
    let even = ifft(&x.iter().step_by(2).cloned().collect::<Vec<_>>());
    let odd = ifft(&x.iter().skip(1).step_by(2).cloned().collect::<Vec<_>>());
    let pi = std::f64::consts::PI;
    let n: f64 = N as f64;
    let T = (0..N/2).map(|k| Complex::new(0.0, 2.0*pi*(k as f64)/n).exp()
                                                * odd[k]).collect::<Vec<_>>();
    let mut v: Vec<Complex<f64>> = vec![Complex::new(0.0, 0.0); N];
    for k in 0..N {
        if k < N / 2 {
            v[k] = (even[k] + T[k]) * 0.5
        }
        else {
            v[k] = (even[k-N/2] - T[k-N/2]) * 0.5
        }
    }
    v
}

fn ceil_2pow(n: usize) -> usize {
    let mut m: usize = 1;
    while m < n {
        m *= 2
    }
    m
}

fn add_zeros(f: &Vec<i64>, N: usize) -> Vec<Complex<f64>> {
    let mut g: Vec<Complex<f64>> = vec![Complex::<f64>::new(0.0, 0.0); N];
    for i in 0..f.len() {
        g[i] = Complex::<f64>::new(f[i] as f64, 0.0)
    }
    g
}

fn fft_polynomial_square(f: &Vec<i64>) -> Vec<i64> {
    let n = f.len() * 2 - 1;
    let N = ceil_2pow(n);
    
    let f_ext = add_zeros(f, N);
    
    let fft_f = fft(&f_ext);
    
    let fft_mul: Vec<Complex<f64>> = (0..N).
                            map(|i| fft_f[i] * fft_f[i]).collect();
    
    let result = ifft(&fft_mul);
    return result.into_iter().map(|c| c.re.round() as i64).
                                                    collect::<Vec<_>>()
}


//////////////////// process ////////////////////

fn read_input() -> Vec<usize> {
    let _N: usize = read();
    let S: Vec<usize> = read_vec();
    S
}

fn frequency(S: Vec<usize>) -> Vec<i64> {
    let max_value = S.iter().map(|&v| v).max().unwrap();
    let mut freq: Vec<i64> = vec![0; max_value+1];
    for v in S.into_iter() {
        freq[v] += 1
    }
    freq
}

fn F(S: Vec<usize>) -> i64 {
    let A = frequency(S);
    let B = fft_polynomial_square(&A);
    let mut counter: i64 = 0;
    for i in 0..A.len() {
        counter += ((B[i*2] - A[i]*A[i]) / 2) * A[i]
    }
    counter
}

fn main() {
    let S = read_input();
    println!("{}", F(S))
}