Counting the number of triplets forming GP series in python

The problem statement is:

You are given an array and you need to find number of triplets of indices(i,j,k) such that the elements at those indices are in geometric progression for a given common ratio, r and i<j<k. For example, arr=[1,4,16,64]. If r=4, we have [1,4,16] and [4,16,64] at indices (0,1,2) and (1,2,3).

I have a solution but I am unable to understand it. Can anyone help me out?

def countTriplets(arr, r):
 count = 0
 
 dict = {}
 dictPairs = {}

 for i in reversed(arr):
    if i*r in dictPairs:
        count += dictPairs[i*r]
    if i*r in dict:
        dictPairs[i] = dictPairs.get(i, 0) + dict[i*r]

    dict[i] = dict.get(i, 0) + 1

 return count

Answer

The way this problem works is that it uses “dict” to keep a running count of each number in the array. Then, it uses “dictPairs” to keep a running count of the pairs for i and (i*r). Lastly, if (i*r) exists in dictPairs, that means there is a triplet, because if (i*r) exists in dictPairs that must mean (i*r*r) exists in dict, and the count is increased by the number of instances of a pair in dictPairs.

Using the example provided and iterating step by step might give a clearer answer:

The first element in the reversed array is 64

• 256 (64*4) is not in dictPairs, nor is it in dict
• The pair (64, 1) is added to dict because it is not in dict yet and thus the default value of 0 is returned from dict.get(i, 0)

The second element in the array is 16

• 64 (16*4) is not in dictPairs, but it is in dict. Thus, (16, 1) is added to dictPairs, because dictPairs does not contain a 16 yet, and dict[64] is equal to one
• The pair (16, 1) is added to dict

The third element in the array is 4

• 16 (4*4) is in dictPairs! This means the count increases by one, because dictPairs[16] is equal to one
• 16 is also in dict, so the pair (16, 1) is added to dictPairs
• The pair (4, 1) is added to dict

The last element in the array is 1

• 4 is in dictPairs, so the count increases by one again, because dictPairs[4] is equal to one
• 4 is also in dict, so the pair (4, 1) is added to dictPairs
• the pair (1, 1) is added to dict

The count is then 2 because there are two triplets.