Update 003. Count Triplets.py

[kartikshastrakar]

add refactored method

[kartikshastrakar]

from collections import defaultdict

def count_triplets(n, r, arr):
arr2 = defaultdict(int)
arr3 = defaultdict(int)
count = 0
for i in arr:
count += arr3[i]
arr3[i * r] += arr2[i]
arr2[i * r] += 1
return count

Example usage

n, r = map(int, input("Enter n and r: ").split())
arr = list(map(int, input("Enter array elements: ").split()))
result = count_triplets(n, r, arr)
print("Number of triplets:", result)

[kartikshastrakar]

add