G - Fine Triplets Editorial
by
physics0523
\(B-A = C-B\) を移項して、 \(2B=A+C\) とします。
ここで、ひとまず \(A,C \in S\) かつ \(A+C=k\) なる整数組 \((A,C)\) の数を数えましょう。
例えば、 \(S = \{1,2,3,5\}\) であるとき、以下のような式を考えます。
\((\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}\)
この式の左辺を展開するとき、左側の \((x^1+x^2+x^3+x^5)\) から \(1\) 項、右側の \((x^1+x^2+x^3+x^5)\) から \(1\) 項を選んで、その積を右辺に加算することになります。
例えば、左側から \(x^5\) 、右側から \(x^2\) を選んだ時、右辺に \(x^{5+2} = x^7\) が加算されます。
これを全体について行うと、順序付きの整数組 \((A,C)\) \((A,C \in S)\) であって \(A+C=k\) なるものの個数が右辺の \(x^k\) につく係数として現れます。
ここで、ふたつの多項式の積は Convolution を使って求めることができます。 (ACL のドキュメント)
注意として、右辺の各項につく係数は最大でも \(10^6\) なので、 \({\rm mod}\ 998244353\) 上で余りを求めてしまっても (ACL の convolution で求まる値をそのまま使っても) 構いません。
この情報を使って求解しましょう。
\(B \in S\) なる全ての \(B\) について、 \(A<B<C\) となるように \(S\) 中から \(A,C\) を選んで \((A,B,C)\) を良い三つ組とする方法が何通りあるか数えて足し合わせます。
\(S = \{1,2,3,5\}\) の例で、 \(B=3\) とするとき \(A+C=2B\) かつ \(A,C \in S\) なる順序付きの整数組 \((A,C)\) は \((1,5),(3,3),(5,1)\) の \(3\) つです。
このうち \((B,B)=(3,3)\) を選んだ場合は \(A=B=C\) となり良い三つ組とはならず、 そうでない組 \((A,C)\) を選択すると良い三つ組を作ることができますが、全ての組が \((C,A)\) と合わせて \(2\) 回ずつカウントされています。
よって、 \(x^{2B}\) の係数から \(1\) を引き \(2\) で割ると \(B\) を固定した際の良い三つ組の数が分かります。
以上を実装することで、この問題に正解できます。時間計算量は \(A_i\) の最大値を \(M\) として \(O(M \log M)\) です。
実装例 (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";
}
実装例 (Python) (PyPy として提出すると TLE する恐れがあります):
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)
posted:
last update:
