CODEWITHSUNDEEP
python
Alexander Soare
Alexander Soare

Reputation: 3267

How to enumerate combinations by tuple indices?

According to the Python docs, combinations follows this sort of pattern:

# combinations('ABCD', 2) --> AB AC AD BC BD CD

I'll be specific about my use case, hopefully without losing generality. I'd like to compute a pairwise metric for a sequence of objects like this:

from itertools import combinations
metrics = [metric(a,b) for (a, b) in combinations(objects, 2)]

Then I want to grab pairs of my objects based on some metric threshold. To do so I'd need to invert this correspondence combinations('ABCD', 2) --> AB AC AD BC BD CD. Ideally, I'd like to avoid writing any logic that does this, and just get it from the source of truth: the combinations function.

I'd like to enumerate like this:

metrics = {(i, j): metric(a,b) for (i, j), (a, b) in enumerate(combinations(objects, 2))}

But of course this doesn't yield a tuple index, just an integer index.

How do I get what I'm looking for?

Upvotes: 0

Views: 71

Answers (3)

no comment
no comment

Reputation: 10616

Redstone Tehnik's is the straight-forward solution, but here's yet another one. Leaving all the hard work to C code, it might be faster:

from itertools import combinations, starmap

metrics = dict(zip(combinations(range(len(objects)), 2),
                   starmap(metric, combinations(objects, 2))))

Little benchmark:

239 ms  240 ms  241 ms  244 ms  253 ms  Alex
197 ms  201 ms  203 ms  206 ms  209 ms  Red
191 ms  191 ms  193 ms  193 ms  194 ms  Dont

Code (Try it online!):

from timeit import repeat
from itertools import combinations, starmap
from operator import sub

def Alex(objects, metric):
    return {(i, j): metric(a,b) for (i, j), (a, b) in zip(combinations(range(len(objects)), 2), combinations(objects, 2))} 

def Red(objects, metric):
    return {(i, j): metric(a,b) for (i, a), (j, b) in combinations(enumerate(objects), 2)}

def Dont(objects, metric):
    return dict(zip(combinations(range(len(objects)), 2), starmap(metric, combinations(objects, 2))))

funcs = Alex, Red, Dont

expect = None
for func in funcs:
    result = func(range(1000), lambda a, b: [a, b])
    if expect is None: expect = result
    print(result == expect, func.__name__)
print()

for _ in range(3):
    for func in funcs:
        times = sorted(repeat(lambda: func(range(1000), lambda a, b: None), number=1))
        print(*('%3d ms ' % (t * 1e3) for t in times), func.__name__)
    print()

Upvotes: 1

Redstone Tehnik
Redstone Tehnik

Reputation: 177

You can enumerate the objects and make combinations of those.

metrics = {(i, j): metric(a,b) for (i, a), (j, b) in combinations(enumerate(objects), 2)}

Upvotes: 2

Alexander Soare
Alexander Soare

Reputation: 3267

Figured this out as I was almost done writing it so decided to finish the question and post an answer. Here it is!

metrics = {(i, j): metric(a,b) for (i, j), (a, b) in zip(combinations(range(len(objects)), 2), combinations(objects, 2))]

So now I can find the largest metric and corresponding key, then use the indices of the key to access my objects sequence.

Upvotes: 0

Related Questions