Official

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: