Last Updated: 29 Nov, 2021
Count the AP
Hard
Asked in companies
Bob and Alice recently learned about the arithmetic progression. Alice found an array ‘A’ of size ‘N’ and wondered how many subsequences of this array of size at least 3 are an arithmetic progression.
A sequence is an arithmetic progression if the difference between one term and the next is a constant.
Bob asked you to help her count the number of such subsequences.
Note: It is guaranteed that the answer fits in 32-bit integer.
Input Format:
The first line of input contains an integer ‘T’, denoting the number of test cases.
The first line of each test case contains one integer ‘N’, denoting the size of the array.
The following line contains an array ‘A’ of ‘N’ spaced integers denoting the array.
Output Format:
For each test case, print a single integer denoting the number of such subsequences.
Note:
You are not required to print the expected output. It has already been taken care of. Just implement the function.
Constraints:
1 <= T <= 5
1 <= N <= 10^3
-10^9 <= A[i] <= 10^9
Time Limit: 1 sec
Approaches
We can generate all possible subsequences using depth-first search and count the number of subsequences that are of size at least 3 and form an arithmetic progression.
Algorithm:
- Initialise ‘ans’ to 0.
- Generate all possible subsequences using DFS
- If the size of the subsequence is less than 3, continue.
- If the subsequence forms an AP, increment ‘ans’ by 1.
- Return ‘ans’
We will be using dynamic programming to optimally calculate the count.
Let, ‘f[i][d]’ denotes the number of arithmetic subsequences that ends with ‘A[i]’ and its common difference is ‘d’ with size >= 2.
Note that, we can append an element ‘A[i]’ at the end of another element ‘A[j]’, only if the difference between the ‘A[j]’ and ‘A[i]’ is equal to the sequence's common difference.
Thus, we can define the state transitions for the element ‘A[i]’ as:
for all j < i
Let diff = ‘A[ i ]’ - ‘A[ j ]’
‘f[ i ][ diff ]’ += ‘f[ j ][ diff ]’ + 1
Here, ‘f[i][d]’ denotes the count of arithmetic subsequences with size >= 2, but we need the count of arithmetic subsequences with size >= 3.
Look at the expression,
‘f[ i ][ diff ]’ += ‘f[ j ][ diff ]’ + 1
Here, ‘f[ j ][ diff ]’ denotes the already existing subseqnce ending at ‘A[j]’, while 1 denote the new subseqnce formed by ‘A[ i ]’ and ‘A[ j ]’. When we are appending new elements to existing subseqnce of size 2, we are forming arithmetic subsequences. So the first part, ‘f[ j ][ dif ]’ is the number of new formed arithmetic subsequences, and can be added to the answer.
Algorithm:
- Initialise ‘ans’ to 0.
- Create a vector of unordered map ‘f’ from long long int to int of size ‘N’.
- Run a loop ‘i’ from 1 to ‘N’ - 1
- Run a loop ‘j’ from 0 to ‘i’ - 1
- Initialise ‘diff’ to ‘A[i]’ - ‘A[j]’.
- Add ‘f[j][diff]’ + 1 to ‘f[i][diff]’.
- Add ‘f[j][diff]’ to ‘ans’.
- Run a loop ‘j’ from 0 to ‘i’ - 1
- Return ‘ans’.