G - Fine Triplets / Time Limit: 3 sec / Memory Lim...

创建于:2025年2月8日

使用 O3 Mini 回答 Chat01

问题

G - Fine Triplets /
Time Limit: 3 sec / Memory Limit: 1024 MB

×
It is prohibited to use generative AI in ongoing AtCoder contests. Please refer to the following rules for more details.

AtCoder Rules against Generative AI - Version 20241206

Score :
600 points

Problem Statement
For integers
A,B,C (
A<B<C ), if they satisfy
B−A=C−B, then
(A,B,C) is called a fine triplet.
You are given a set of
N distinct positive integers
S={S
1

,S
2

,…,S
N

}. Find the number of fine triplets
(A,B,C) with
A,B,C∈S.

Constraints
All input values are integers.
1≤N≤10
6

1≤S
i

≤10
6

The elements of
S are distinct.
Input
The input is given from Standard Input in the following format:

N
S
1

S
2


S
N

Output
Print the number of fine triplets as an integer.

Sample Input 1
Copy
5
8 3 1 5 2
Sample Output 1
Copy
3
Here,
S={8,3,1,5,2}. The fine triplets to be counted are the following three:

(1,2,3)
(1,3,5)
(2,5,8)
Sample Input 2
Copy
7
300000 100000 499998 499999 200000 400000 500000
Sample Output 2
Copy
5
Sample Input 3
Copy
10
13 1 16 15 12 4 7 10 2 19
Sample Output 3
Copy
10
Language
Source Code

at most 512 KiB
Your source code will be saved as Main.extension.
Please give correct c++ code all tests pass please and also give FULL all test pass code please and look constraints and using constratints give optimzed code but all tests pass code please take care of this...!!
all tests paas please

思考

Analyzing fine triplets

So, the task is about counting distinct triplets (A, B, C) where A < B < C and they form an arithmetic progression, meaning B - A = C - B. This simplifies to the condition (C + A) = 2 * B. It's based on counting 3-element arithmetic progressions from a set S of distinct integers. Though the set isn't sorted, I can sort it and then look for triplets that meet this condition. This seems like a classic arithmetic progression problem!

Counting fine triplets efficiently

To count arithmetic triplets, I can iterate through all possible candidates for B and check if there are two numbers A and C such that A < B < C and A + C = 2*B. The approach involves sorting S and checking for each B using a hash set, but iterating over each pair is slow (O(N^2)).

A more efficient idea is to iterate over A and check if B = A + d and C = A + 2d. However, iterating over d could be too large. I need something faster, maybe using two-pointer technique if S is sorted.

Optimizing arithmetic progressions counting

The problem is to count triplets (A, B, C) in an arithmetic progression from S. I thought of iterating through each B and checking for pairs (A, C) such that A = B - d and C = B + d. But iterating over all differences d seems inefficient.

A better approach involves creating a frequency array for elements of S, and then searching for matching A and C for each potential B. Even so, it requires efficient summing and quicker membership tests—still a challenge at N = 1e6. I need a clever approach to improve performance.

Using FFT for counting triplets

The formula for counting arithmetic progressions connects to the convolution of an indicator function of S. By computing the convolution f * f, we can count pairs (A, C) such that A + C = 2B for each B. Each valid pair is counted twice, so we divide by 2. Calculating this using FFT helps with efficiency, but the precision and handling of large numbers could be tricky for dense S. Still, it’s much faster than naïve iteration, and O(M log M) complexity is feasible for large M (up to 1e6).

Implementing FFT for counting fine triplets

The goal is to count fine triplets using FFT. I’ll start by implementing a convolution using the FFT method. The polynomial f will be created as an array where the values of S are set to 1. Then, I will compute the FFT of this polynomial, square the result, and apply the inverse FFT to get the number of valid pairs. Finally, I'll sum the answers for each b in S, checking if 2b exists in the convolution array. Careful handling of precision and boundary conditions (like 2b) is key here.

Fixing convolution overcounts

It looks like the convolution approach is slightly overcounting, especially when considering cases like (b, b). For a valid triplet (A, B, C), with A < B < C, when computing f*f(2b), the pair (b, b) is counted, which results in an extra term because we need A and C to be distinct.

To address this, the formula needs to subtract 1 from the convolution count when b is in S, ensuring distinct pairs are counted correctly. I'll adjust for this and recheck computations for all b in the set S.

Reasoned for 47 seconds

回答

We can solve this problem by “counting” each valid arithmetic progression in SS. Recall that a triple (A,B,C)(A,B,C) (with A<B<CA<B<C) is a fine triplet if and only if

BA=CBA+C=2B.B - A = C - B \quad \Longleftrightarrow \quad A + C = 2B.

One neat idea is to “fix” the middle element BB (which must lie in SS) and then count how many pairs (A,C)(A,C) (with A,CSA,C\in S) satisfy

A+C=2B.A + C = 2B.

Notice that if (A,C)(A,C) is a solution (with ACA\ne C) then automatically one of AA or CC is less than BB and the other is greater than BB (since if both were on one side the sum would be either too small or too large).

A standard observation is that if we define an indicator function

f(x)={1if xS,0otherwise,f(x)=\begin{cases}1 &\text{if }x\in S,\\0&\text{otherwise},\end{cases}

then for any integer kk the convolution

(ff)(k)=x+y=kf(x)f(y)(f*f)(k)=\sum_{x+y=k} f(x)f(y)

counts the number of pairs (x,y)(x,y) with x,ySx,y\in S and x+y=kx+y=k.

Now, note that when BSB\in S we have:

  • One of the terms in the sum is when x=y=Bx=y=B (since B+B=2BB+B=2B); however, (B,B)(B,B) is not allowed because we require the three numbers to be distinct (and A<B<CA<B<C).
  • All other pairs come in symmetric pairs: if (A,C)(A,C) is a solution with ACA\neq C then both (A,C)(A,C) and (C,A)(C,A) appear in the sum.

Thus, if BSB\in S then

(ff)(2B)=1+2(# of valid pairs (A,C) with A+C=2B).(f*f)(2B)=1+2\cdot(\#\text{ of valid pairs }(A,C)\text{ with } A+C=2B).

In other words, the number of valid pairs (which give a fine triplet with middle BB) is

(ff)(2B)12.\frac{(f*f)(2B)-1}{2}.

Thus the final answer is:

ans=BS(ff)(2B)12.\text{ans}=\sum_{B\in S} \frac{(f*f)(2B)-1}{2}.

The challenge is to compute the convolution fff*f quickly. Since all SiS_i are positive integers 106\le 10^6 and N106N\le 10^6, we can define an array of length (say) m=106+1m=10^6+1 where

f(x)=1 if xS,0 otherwise.f(x)=1 \text{ if } x\in S,\quad 0 \text{ otherwise.}

Then the convolution fff*f can be computed via Fast Fourier Transform (FFT) in O(mlogm)O(m\log m) time. (The resulting convolution will have length up to roughly 21062\cdot10^6.)

The code below implements the FFT (using the standard iterative method) and then computes the convolution. Finally, for each BSB\in S we add

round((ff)(2B))12\frac{\text{round}((f*f)(2B))-1}{2}

to the answer. (We use llround to convert from the floating‐point result of the FFT back to an integer.)

Below is the full optimized C++ code:


cpp
#include <bits/stdc++.h> using namespace std; typedef complex<double> cd; // Iterative FFT implementation. void fft(vector<cd> & a, bool invert) { int n = a.size(); for (int i = 1, j = 0; i < n; i++) { int bit = n >> 1; for(; j & bit; bit >>= 1) j -= bit; j += bit; if(i < j) swap(a[i], a[j]); } for (int len = 2; len <= n; len <<= 1) { double ang = 2 * M_PI / len * (invert ? -1 : 1); cd wlen(cos(ang), sin(ang)); for (int i = 0; i < n; i += len) { cd w(1); for (int j = 0; j < len/2; j++) { cd u = a[i+j], v = a[i+j+len/2] * w; a[i+j] = u + v; a[i+j+len/2] = u - v; w *= wlen; } } } if(invert) { for (int i = 0; i < n; i++) a[i] /= n; } } int main(){ ios::sync_with_stdio(false); cin.tie(nullptr); int N; cin >> N; vector<int> S(N); int maxVal = 0; for (int i = 0; i < N; i++){ cin >> S[i]; if(S[i] > maxVal) maxVal = S[i]; } // Our f(x) array will be of size m = maxVal + 1. int m = maxVal + 1; // Build the indicator function f. vector<cd> poly(m, 0); for (int x : S) { poly[x] = 1; } // We need to compute the convolution f*f. // The required size is at least (2*m - 1). int n = 1; while(n < 2 * m - 1) n <<= 1; poly.resize(n); // Perform FFT on poly. fft(poly, false); for (int i = 0; i < n; i++){ poly[i] *= poly[i]; } fft(poly, true); // Now poly[k] (rounded) is equal to (f*f)(k). // For each B in S, the number of fine triplets with B as the middle is ((poly[2*B] - 1) / 2). long long ans = 0; for (int b : S) { int idx = 2 * b; if(idx < (int)poly.size()){ long long convVal = (long long)llround(poly[idx].real()); // Subtract the self-pair (b,b) convVal -= 1; convVal /= 2; if(convVal > 0) ans += convVal; } } cout << ans << "\n"; return 0; }

Explanation

  1. Input Reading and Setup:

    • We read NN and the set SS.
    • We determine the maximum value in SS (denoted as maxVal) and set the size m=maxVal+1m = \texttt{maxVal}+1.
    • We create an array poly (of type vector<cd>) such that poly[x] = 1 if xSx \in S and 0 otherwise.
  2. FFT Setup:

    • We resize the polynomial to length n, the next power of 2 at least 2m12*m - 1.
    • We perform FFT on poly, square each value (this corresponds to convolution in the frequency domain) and then perform the inverse FFT to obtain the convolution fff*f.
  3. Counting the Fine Triplets:

    • For every BSB\in S we consider index 2B2B in the convolution array. Since the pair (B,B)(B,B) is always present (and must be removed), we subtract 1 and then divide by 2.
    • We sum these counts for all BSB\in S and print the answer.

This code is optimized to run under the given constraints and should pass all test cases.

Happy coding!

分享这个问答