Count triplets

DSA/Python/Hashing/Count Triplets/Problem Statement

@@ -0,0 +1,17 @@
+You are given an array and you need to find number of tripets 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 .
+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 n and r, the size of arr and the common ratio.
+The next line contains n space-seperated integers arr[i].

DSA/Python/Hashing/Count Triplets/Solution

@@ -0,0 +1,18 @@
+from collections import Counter
+
+def countTriplets(arr, r):
+    a = Counter(arr)
+    b = Counter()
+    s = 0
+    for i in arr:
+        j = i//r
+        k = i*r
+        a[i]-=1
+        if b[j] and a[k] and i%r == 0:
+            s+=b[j]*a[k]
+        b[i]+=1
+    return s
+
+n,r = map(int,input().split())
+arr = list(map(int,input().split()))
+print(countTriplets(arr, r))