Search code examples
pythonalgorithmsortingpython-itertools

Optimize python algorithm


I have three sorted lists, examplewise

a, b, c = [10,9,8], [9,8,7], [13,5,1]

I want to get all the combinations x, y, z where x in a, y in b and z in c and 1/x + 1/y + 1/z < 1 in the fastest time possible. I've been trying some different approaches,

for x, y, z in product(a, b, c):
    if predicative(x,y,z):
        yield (x, y, z)

Obviously, this takes too long time considering I'm checking everything, and the lists a, b, c are already sorted. I have tried sorting product(a,b,c) on the sum, but that is reaaally slow, as it uses all the products. My initial plan with having a, b and c sorted is so I could break out of the loop as soon as one fails. Any ideas?

Thanks.


Solution

  • A simple solution that could speed it up a bit would be to store 1/z for each z in c in a list, and for each pair of x,y in a,b - use binary search for the highest 1/z (in the auxillary list) such that 1/x + 1/y + 1/z < 1 - this will effectively trim many search from the 3rd lists and get you some speed up.
    This will reduce time complexity to O(log(n)*n^2 + Y), where Y is the output size (number of triplets produced).

    Note however, that since the output size itself is O(n^3) (consider 3 lists with all elements > 3.333) - you cannot avoid the worst case slow time, since you might be needing to generate n^3 triplets.

    If you want only the count however, the approach suggested can easily find it in O(n^2*logn).