Given some integers, find triplets where the product of two numbers equals the third number I am trying to solve a coding challenge:


  Given N integers A1, A2, …, AN, count the number of triplets (x, y, z) (with 1 ≤ x &lt; y &lt; z ≤ N) such that at least one of the following is true:
  
  Ax = Ay × Az, and/or
  Ay = Ax × Az, and/or
  Az = Ax × Ay
  
  Sample case 1
  
  5 2 4 6 3 1
  
  In Sample Case #1, the only triplet satisfying the condition given in the problem statement is (2, 4, 5). The triplet is valid since the second, fourth, and fifth integers are 2, 6, and 3, and 2 × 3 = 6. thus the answer here is 1.
  
  Sample case 2
  
  2 4 8 16 32 64
  
  The six triplets satisfying the condition given in the problem statement are: (1, 2, 3), (1, 3, 4), (1, 4, 5), (1, 5, 6), (2, 3, 5), (2, 4, 6). so the answer here is 6.


My Code in python:

import itertools
count=0
for t in itertools.combinations(l,3):
    if t[0]*t[1]==t[2] or t[1]*t[2]==t[0] or t[0]*t[2]==t[1]:
        count+=1
print(count)


This is the naive way of generating all possible 3 length combinations and checking for the condition. This works fine for smaller input but when the inout size increase time complexity increases. I am assuming for an example that has 1,2,3,6,8 the combinations generated are (2,3,6),(2,3,8) 2,3,6 satisfy the condition so the checking for 2,3,8 is unnecessary and can be avoided. How can I modify my code to take advantage of this observation ?
Your combinations(…, 3) loop makes your algorithm O(N3).

It's easy to improve it to be O(N2).  The question is, essentially: for every pair of entries, how many occurrences of their product are in the list?  So, make an indexes data structure to help you find, in O(1) time, where the product might be located.

from collections import defaultdict
from itertools import combinations

a = [int(ai) for ai in input('Input: ').split()]

indexes = defaultdict(set)
for i, ai in enumerate(a):
    indexes[ai].add(i)

triplets = set()
for x, y in combinations(range(len(a)), 2):
    for z in indexes[a[x] * a[y]].difference([x, y]):
        triplets.add(tuple(sorted((x, y, z))))

print(len(triplets))


Here, I've chosen to stick closer to the notation used in the challenge itself, with a being the list, and x, y, z as the indexes of entries (but 0-based rather than 1-based).