We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
  1. Count Triplets
  2. Discussions

Count Triplets

  • nikolai_miroshn1
     + 0 comments

    I got test 6 and 10 wrong. Any help will be appreciated!

    Here is my code:

    #!/bin/python3

import math
import os
import random
import re
import sys

# Complete the countTriplets function below.
def countTriplets(arr, r):
    num_counts = {}
    triplets = 0
    # creating the hashmap
    for i in arr:
        if i in num_counts:
            num_counts[i] += 1
        else:
            num_counts[i] =1
    # counting tripplets
    

    for key in num_counts:
        # if ar and ar^2 is not in the map skip
        a, ar, arr = 0, 0, 0
        if key*r not in num_counts or key*r**2 not in num_counts:
            continue
        if r == 1:
            a = num_counts[key]
            triplets += math.factorial(a)//(math.factorial(a-3)*6)
        else:
            a = num_counts[key]
            ar = num_counts[key*r]
            arr = num_counts[key*r**2]
            triplets += a*ar*arr
        
    return triplets

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()