A possible solution to count triplets HackerRank problem (https://www.hackerrank.com/challenges/count-triplets-1)

count_triplets.py

!/bin/python3
import collections
import math
import os
import random
import re
import sys


def countPairs(arr, r):
  pair_count = 0
  element_freq_map = collections.defaultdict(int)
  for i in range(len(arr)):
    if arr[i]/r in element_freq_map:
      pair_count += element_freq_map[arr[i]/r]
    element_freq_map[arr[i]] += 1
  return pair_count


def countTripletsBruteForce(arr, r):
  count_trips = 0
  ## Brute force - works but is too slow
  # This is basically O(N^3)
  for i in range(len(arr)):
    for j in range(i+1, len(arr)):
      for k in range(j+1, len(arr)):
        if (arr[j] / arr[i] == r
            and arr[k] / arr[j] == r):
              count_trips += 1
   return count_trips


# Complete the countTriplets function below.
def countTriplets(arr, r):
  count_trips = 0
  # O(N)-ish? (actually editorial suggests it is O(NlgN))
  element_freq_map = collections.defaultdict(int)
  pair_freq_map = collections.defaultdict(int)
  for i in range(len(arr)):
    if arr[i]/r in pair_freq_map: 
      # If arr[i]/r is in the pair frequency map,
      # then it means there is a triple ending with
      # this element, so increment the counter by
      # the number of pairs that end with arr[i]/r
      count_trips += pair_freq_map[arr[i]/r]
    if arr[i]/r in element_freq_map:
      # If arr[i]/r is in the element freq map,
      # then it means there is/are pair(s) ending with
      # this element, arr[i]. Store and increment the
      # pair frequency map
      pair_freq_map[arr[i]] += element_freq_map[arr[i]/r]
    element_freq_map[arr[i]] += 1
  return count_trips


if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    nr = input().rstrip().split()

    n = int(nr[0])

    r = int(nr[1])

    arr = list(map(int, input().rstrip().split()))

    ans = countTriplets(arr, r)

    fptr.write(str(ans) + '\n')

    fptr.close()