fastest algorithm count number of 3 length AP in array - c++

Question

I want to solve this CodeChef challenge:
Suppose We are given an array A of N(of range 100,000) elements. We are to find the count of all pairs of 3 such elements 1<=Ai,Aj,Ak<=30,000 such that
Aj-Ai = Ak- Aj and i < j < k
In other words Ai,Aj,Ak are in Arithmetic Progression. For instance for Array :
9 4 2 3 6 10 3 3 10
so The AP are:
{2,6,10},{9,6,3},{9,6,3},{3,3,3},{2,6,10}
So the required answer is 5.
My Approach
What I tried is take 30,000 long arrays named past and right. Initially right contains the count of each 1-30,000 element.
If we are at ith position past stores the count of array value before i and right stores the count of array after i. I simply loop for all possible common difference in the array. Here is the code :
right[arr[1]]--;
for(i=2;i<=n-1;i++)
{
past[arr[i-1]]++;
right[arr[i]]--;
k=30000 - arr[i];
if(arr[i] <= 15000)
k=arr[i];
for(d=1;d<=k;d++)
{
ans+= right[arr[i] + d]*past[arr[i]-d] + past[arr[i] + d]*right[arr[i]-d];
}
ans+=past[arr[i]]*right[arr[i]];
}
But this gets me Time Limit Exceeded. Please help with a better algorithm.

Answer 1

You can greatly cut execution time if you make a first pass over the list and only extract number pairs that it is possible to have an 3 term AP between (difference is 0 mod 2). And then iterating between such pairs.
Pseudo C++-y code:
// Contains information about each beginning point
struct BeginNode {
int value;
size_t offset;
SortedList<EndNode> ends; //sorted by EndNode.value
};
// Contains information about each class of end point
struct EndNode {
int value;
List<size_t> offsets; // will be sorted without effort due to how we collect offsets
};
struct Result {
size_t begin;
size_t middle;
size_t end;
};
SortedList<BeginNode> nodeList;
foreach (auto i : baseList) {
BeginNode begin;
node.value = i;
node.offset = i's offset; //you'll need to use old school for (i=0;etc;i++) with this
// baseList is the list between begin and end-2 (inclusive)
foreach (auto j : restList) {
// restList is the list between iterator i+2 and end (inclusive)
// we do not need to consider i+1, because not enough space for AP
if ((i-j)%2 == 0) { //if it's possible to have a 3 term AP between these two nodes
size_t listOffset = binarySearch(begin.ends);
if (listOffset is valid) {
begin.ends[listOffset].offsets.push_back(offsets);
} else {
EndNode end;
end.value = j;
end.offsets.push_back(j's offset);
begin.ends.sorted_insert(end);
}
}
}
if (begin has shit in it) {
nodeList.sorted_insert(begin);
}
}
// Collection done, now iterate over collection
List<Result> res;
foreach (auto node : nodeList) {
foreach (auto endNode : node.ends) {
foreach (value : sublist from node.offset until endNode.offsets.last()) {
if (value == average(node.value, endNode.value)) {
// binary_search here to determine how many offsets in "endNode.offsets" "value's offset" is less than.
do this that many times:
res.push_back({node.value, value, endNode.value});
}
}
}
}
return res;

Answer 2

Here's a simple C version of the solution that takes advantage of the Ai + Ak must be even test:
#include <stdio.h>
static int arr[] = {9, 4, 2, 3, 6, 10, 3, 3, 10};
int main ()
{
int i, j, k;
int sz = sizeof(arr)/sizeof(arr[0]);
int count = 0;
for (i = 0; i < sz - 2; i++)
{
for (k = i + 2; k < sz; k++)
{
int ik = arr[i] + arr[k];
int ikdb2 = ik / 2;
if ((ikdb2 * 2) == ik) // if ik is even
{
for (j = i + 1; j < k; j++)
{
if (arr[j] == ikdb2)
{
count += 1;
printf("{%d, %d, %d}\n", arr[i], arr[j], arr[k]);
}
}
}
}
}
printf("Count is: %d\n", count);
}
and the console dribble:
tmp e$ cc -o triples triples.c
tmp e$ ./triples
{9, 6, 3}
{9, 6, 3}
{2, 6, 10}
{2, 6, 10}
{3, 3, 3}
Count is: 5
tmp e$
This more complicated version keeps a list of Aj indexed by value to go from n-cubed to n-squared (kinda).
#include <stdio.h>
#include <stdint.h>
static uint32_t arr[] = {9, 4, 2, 3, 6, 10, 3, 3, 10};
#define MAX_VALUE 100000u
#define MAX_ASIZE 30000u
static uint16_t index[MAX_VALUE+1];
static uint16_t list[MAX_ASIZE+1];
static inline void remove_from_index (int subscript)
{
list[subscript] = 0u; // it is guaranteed to be the last element
uint32_t value = arr[subscript];
if (value <= MAX_VALUE && subscript == index[value])
{
index[value] = 0u; // list now empty
}
}
static inline void add_to_index (int subscript)
{
uint32_t value = arr[subscript];
if (value <= MAX_VALUE)
{
list[subscript] = index[value]; // cons
index[value] = subscript;
}
}
int main ()
{
int i, k;
int sz = sizeof(arr)/sizeof(arr[0]);
int count = 0;
for (i = 0; i < sz - 2; i++)
{
for (k = i; k < sz; k++) remove_from_index(k);
for (k = i + 2; k < sz; k++)
{
uint32_t ik = arr[i] + arr[k];
uint32_t ikdb2 = ik / 2;
add_to_index(k-1); // A(k-1) is now a legal middle value
if ((ikdb2 * 2) == ik) // if ik is even
{
uint16_t rover = index[ikdb2];
while (rover != 0u)
{
count += 1;
printf("{%d, %d, %d}\n", arr[i], arr[rover], arr[k]);
rover = list[rover];
}
}
}
}
printf("Count is: %d\n", count);
}
and the dribble:
tmp e$ cc -o triples triples.c
tmp e$ ./triples
{9, 6, 3}
{9, 6, 3}
{2, 6, 10}
{2, 6, 10}
{3, 3, 3}
Count is: 5
tmp e$