Speed up counting the number of A[i] * A[j] * A[k] = A[l] where i < j < k < l

Question

I recently got onto the following problem: we consider the following array:

A = [2, 3, 6, 1, 6, 4, 12, 24]

we need to count the number of times these two conditions are satisfied within the array:

A[i] * A[j] * A[k] = A[l] so that 0 <= i < j < k < l < len(A)

for this example the result should be 8. example of satisfied conditions within the array:

2 * 3 * 1 = 6
2 * 6 * 1 = 12
6 * 1 * 4 = 24
3 * 1 * 4 = 12

The straightforward solution I created using python:

A = [2, 3, 6, 1, 6, 4, 12, 24]
result = 0
for i in range(len(A)):
    for j in range(i + 1, len(A)):
        for k in range(j + 1, len(A)):
            for l in range(k + 1, len(A)):
                if A[i] * A[j] * A[k] == A[l]:
                    result += 1
print(result)

I need to find a way to speed up the program using dynamic, maybe memoization or pre computing.

result = 0
for i in range(A):
    for j in range(i + 1, A):
        for k in range(j + 1, A):
             #TODO
print(result)

I was thinking of creating a dictionary that contains a set of dictionaries for each number, to indicate the number and its position, example:

speed_up = {
    6: {
        2: True
        4: True
    }, 
    1: {
       3: True
    },
    ...
}

then we check like the following:

result = 0
for i in range(A):
    for j in range(i + 1, A):
        for k in range(j + 1, A):
            x = A[i] * A[j] * A[k]
            if x in speed_up:
                result += len([z[0] for z in speed_up[x].items() if z[0] > k])

This way we will count all occurrences of the same number after the index k at once, we will not have to go through all array.

Please let me know if my optimisation is flawed, and if there's better optimisation using dynamic programming techniques such as memoization or pre computing.

Answer 1

Here is a simpler version of your algorithm, using a simpler dictionary structure.

result = 0
cnt = {}
for elem in A:
    if elem not in cnt:
        cnt[elem] = 0
    cnt[elem] += 1

for k in range(len(A)):
    cnt[A[k]] -= 1  # uncount A[k]
    for j in range(k):
        for i in range(j):
            x = A[i] * A[j] * A[k]
            if x in cnt and cnt[x] > 0:
                result += cnt[x]

Notice that through the loop structure we ensure i < j < k and that A[0,...,k] are always erased from the counts in the cnt dictionary. This runs in $O(n^3)$ (assuming $O(1)$ hashmap)