Miraacle's blog

Finding number of arithmetic progressions of length 3 in array.
By Miraacle, history, 20 months ago, In English

Hey, I've been trying to solve the following problem. Given an array of $$$n \leq 2*10^5$$$ distinct positive integers $$$a_1, a_2, \dots, a_n$$$ find the number of triplets ($$$i$$$, $$$j$$$, $$$k$$$) and $$$i < j < k$$$ such that $$$a_j - a_i = a_k - a_j$$$ (so $$$a_i$$$, $$$a_j$$$, $$$a_k$$$ form an arithmetic progression). Additionally, the test cases were made such that the answer is not greater than $$$10^6$$$.

I could only come up with a $$$O(n^2)$$$ solution and have no idea how to go further (maybe some FFT/Convolutions?). Any hints/help would be appreciated, thanks in advance!

  • Vote: I like it
  • +11
  • Vote: I do not like it

Comments Comments (11)
| Write comment?
»
20 months ago, # |
  Vote: I like it 0 Vote: I do not like it

Why would you need FFT/convolutions? An $$$O(nlogn)$$$ solution is easy to think of & implement.

Sort the array first. Then enumerate through all $$$j$$$ s. Therefore the formula is transformed into $$$2\times a_j=a_i+a_k$$$. The left part is determined; the right can be solved using binary search.

→ Reply
  • »
    »
    20 months ago, # ^ |
      Vote: I like it +10 Vote: I do not like it

    How would you solve it with binary search? Also, remember that the order of given integers matters (for $$$1$$$, $$$2$$$, $$$3$$$ answer is $$$1$$$ while for $$$2$$$, $$$3$$$, $$$1$$$ answer is $$$0$$$).

    → Reply
»
20 months ago, # |
← Rev. 2   Vote: I like it 0 Vote: I do not like it

https://codeforces.com/problemset/problem/452/F

Hash + segment tree

→ Reply
  • »
    »
    20 months ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    It's not the same problem as in my problem you need to find number of different triples, while in the problem you've posted you need to check whether there exists any triple.

    → Reply
    • »
      »
      »
      20 months ago, # ^ |
      ← Rev. 3   Vote: I like it 0 Vote: I do not like it

      I'm sorry I can't help you.:(

      When the sequence is arranged, the difficulty of determining the existence is * 2700, and the difficulty of counting may be greater than 3000

      → Reply
»
20 months ago, # |
  Vote: I like it +6 Vote: I do not like it

You can find this problem on CodeChef. https://www.codechef.com/problems/COUNTARI

→ Reply
»
20 months ago, # |
  Vote: I like it 0 Vote: I do not like it

Excuse me, but I think this is a variant of 3SUM? If this is true, then there would be no easy way to reduce it to under $$$O(n^2)$$$ time. Can we know the range of the array's elements? It may give a hint on whether FFT works or not.

→ Reply
  • »
    »
    20 months ago, # ^ |
    ← Rev. 3   Vote: I like it 0 Vote: I do not like it

    Yes, the only way to do it in O(nlogn) is using FFT, and that can be optimized a little more with blocks, like than we find the number of primes in the interval.

    → Reply
  • »
    »
    20 months ago, # ^ |
      Vote: I like it 0 Vote: I do not like it

    The input array $$$a_1, a_2, a_3, \dots, a_n$$$ is a permutation of $$$1, 2, \dots, n$$$ so all elements are distinct and $$$1 \leq a_i \leq n$$$.

    → Reply
    • »
      »
      »
      20 months ago, # ^ |
        Vote: I like it 0 Vote: I do not like it

      Now I am thinking if this problem may be some usage of prefix sum + suffix sum + DP, but I am not very sure on how. Can anyone help me on this approach?

      → Reply
»
18 months ago, # |
  Vote: I like it +5 Vote: I do not like it

There is a general approach using Fourier analysis.

Consider trilinear functional

$$$AP(f, g, h) = \mathbf{E}_{n,d\in \mathbf{Z}_N} f(n) g(n+d) h(n+2d)$$$

Let these functions be indicator functions $1_A$, $$$1_B$$$, $$$1_C$$$ for some sets $$$A,B$$$ and $$$C$$$ respectively. Then it's clear how to express the number of arithmetic progressions in $$$A\times B\times C$$$ from $$$AP(1_A,1_B,1_C)$$$.

Then apply Fourier transform to them

$$$ \hat 1_A(u) = \mathbf{E}_{x\in \mathbf{Z}_N} 1_A(x) e(-xu), \quad 1_A(x) = \sum_{u\in \mathbf{Z}_N} \hat 1_A(u) e(xu) $$$

Substitute back

$$$ AP(1_A, 1_B, 1_C) = \mathbf{E}_{n,d\in \mathbf{Z}_N} \left(\sum_{u\in \mathbf{Z}_N} \hat 1_A(u) e(nu)\right) \left(\sum_{v\in \mathbf{Z}_N} \hat 1_B(u) e((n+d)v)\right) \left(\sum_{w\in \mathbf{Z}_N} \hat 1_C(w) e((n+2d)w)\right) $$$

and do some algebra

$$$ AP(1_A, 1_B, 1_C) = \mathbf{E}_{n,d\in \mathbf{Z}_N} \sum_{u,v,w} \hat 1_A(u) \hat 1_B(v) \hat 1_C(w) e(nu+(n+d)v+(n+2d)w) $$$
$$$ AP(1_A, 1_B, 1_C) = \sum_{u,v,w} \hat 1_A(u) \hat 1_B(v) \hat 1_C(w) \mathbf{E}_{n} e((u+v+w)n) \mathbf{E}_d e((v+2w)d) $$$

From properties of characters, we know that $$$\mathbf{E}_x e^{kx} = [k=0]$$$ (Iverson bracket notation), hence

$$$ AP(1_A, 1_B, 1_C) = \sum_{a\in \mathbf{Z}_N} \hat 1_A(a) \hat 1_B(-2a) \hat 1_C(a) $$$
→ Reply