Count Triplets (Hackerrank)

You are given an array and you need to find number of tripets of indices  such that the elements at those indices are in geometric progression for a given common ratio  and .

Example
 

There are  and  at indices  and . Return .

Function Description

Complete the countTriplets function in the editor below.

countTriplets has the following parameter(s):

• int arr[n]: an array of integers
• int r: the common ratio

Returns

• int: the number of triplets

Input Format

The first line contains two space-separated integers  and , the size of  and the common ratio.
The next line contains  space-seperated integers .

Constraints

Sample Input 0

4 2
1 2 2 4

Sample Output 0

2

Explanation 0

There are  triplets in satisfying our criteria, whose indices are  and 

Sample Input 1

6 3
1 3 9 9 27 81

Sample Output 1

6

Explanation 1

The triplets satisfying are index , , , ,  and .

Sample Input 2

5 5
1 5 5 25 125

Sample Output 2

4

Explanation 2

The triplets satisfying are index , , , .

Solution :

YouTube Detailed Explanation - Count Triplets (Click)

Approach - Mapping:

The problem is very simple if we divide it into small parts. Like,

1. Take two map beforeMap and afterMap to store left elements and right elements respectively
2. At the beginning afterMap will be full and beforeMap will be empty 
3. Loop through each element and check for 3 conditions : (1) current%r==0 (2) current/r present in beforeMap (3) current*r present in afterMap. If all three conditions are true,

4. Then only count+=beforeFreq(current/r)*afterFreq(current*r)

Code: 

// Complete the countTriplets function below.
    private static long countTriplets(List arr, long r) {
        Map beforeMap = new HashMap<>();
        Map afterMap = new HashMap<>();
        long count = 0;
        for (long num : arr) {
            Long freq = afterMap.get(num);
            afterMap.put(num,freq==null?1:freq+1);
        }
        for(long num : arr) {
            // reduce the count of current element in afterMap
            afterMap.put(num,afterMap.get(num)-1);
            // checking for 1, 2, 3 conditions
            if(num%r==0 && (beforeMap.containsKey(num/r)&&beforeMap.get(num/r)>0) && (afterMap.containsKey(num*r)&&afterMap.get(num*r)>0)) {
                count+=beforeMap.get(num/r)*afterMap.get(num*r);
            }
            if(beforeMap.containsKey(num)) {
                beforeMap.put(num,beforeMap.get(num)+1);
            }
            else {
                beforeMap.put(num,1L);
            }
        }
        return count;
    }

Space is O(n) and time is O(n).