G - Fine Triplets 解説 by en_translator
Transposing the terms in \(B-A = C-B\), we get \(2B=A+C\).
First, let us count the number of integer pairs \((A,C)\) with \(A,C \in S\) and \(A+C=k\).
For example, if \(S = \{1,2,3,5\}\), consider the following expression.
\((\sum_{k \in S} x^k)^2 = (x^1 + x^2+x^3+x^5)(x^1+x^2+x^3+x^5) = x^2 + 2 x^3 + 3 x^4 + 2 x^5 + 3 x^6 + 2 x^7 + 2 x^8 + x^{10}\)
In the expansion of the left hand side of the equation above, we choose one term from \((x^1+x^2+x^3+x^5)\) to the left and one term from \((x^1+x^2+x^3+x^5)\) to the right, and add their product to the right hand side.
For example, when we choose \(x^5\) from the left and \(x^2\) from the right, we add \(x^{5+2} = x^7\) to the right hand side.
After performing this for all pairs, the number of ordered pairs \((A,C)\) of integers \((A,C \in S)\) with \(A+C=k\) occurs as the coefficient of \(x^k\) in the right hand side.
Here, the product of two polynomials can be found as their Convolution. (ACL (AtCoder Library) Document)
Note that the coefficient of the each term in the right hand side is no greater than \(10^6\), so one may find the convolution modulo \(998244353\) (which we can directly obtain with the convolution
function in ACL).
We will use this information to find the answer.
For each \(B \in S\), we count how many ways to choose \(A\) and \(C\) from \(S\) so that \(A<B<C\) and \((A,B,C)\) forms a good triple, and sum up the counts.
In the example for \(S = \{1,2,3,5\}\), if we have \(B=3\), there are three ordered pairs \((A,C)\) with \(A+C=2B\) and \(A,C \in S\).
Among them, picking \((A,C)=(3,3)\) makes \(A=B=C\) which is not applicable for a good triple. The other pairs form good triples, but every pair is counted twice along with \((C,A)\).
Thus, the sought number of good triples for the fixed \(B\) can be obtained as the coefficient of \(x^{2B}\) subtracted by \(1\) then divided by \(2\).
You can get the correct answer by the algorithm above. The time complexity is \(O(M \log M)\), where \(A\) denotes the maximum value of \(A_i\).
Sample code (C++):
#include<bits/stdc++.h>
#include<atcoder/convolution>
using namespace std;
using namespace atcoder;
int main(){
int n;
cin >> n;
vector<int> p(1000005,0);
for(int i=0;i<n;i++){
int a;
cin >> a;
p[a]=1;
}
long long res=0;
auto c=convolution(p,p);
for(int i=0;i<1000005;i++){
if(p[i]==1){
res+=(c[2*i]-1)/2;
}
}
cout << res << "\n";
}
Sample code (Python) (submitting it as PyPy may result in TLE (Time Limit Exceeded)):
import numpy.fft as F
M = 1000005
N,*A = map(int,open(0).read().split())
f = [0]*(M+1)*2
for a in A: f[a]=1
g = F.ifft(F.fft(f,1<<21)**2)
res = 0
for v in A:
res += int(g[2*v]+0.5)
res -= 1
res //= 2
print(res)
投稿日時:
最終更新: