Itertools / nth product

I saw that the itertools recipe collection added a recipe for the nth_combination and wondered if it would be appropriate to add nth_product as well?

Recipe:

import math

def nth_product(idx, *args):
    """returns the nth product of the given iterables.

    Args:
        idx (int): the index.
        *args: the iterables.
    """
    if not isinstance(idx, int):
        raise TypeError(f"Expected int, not {type(idx)}")
    total = math.prod([len(a) for a in args])
    if idx < 0:
        idx += total
    if idx < 0 or idx >= total:
        raise IndexError(f"Index {idx} out of range")
    
    elements = ()
    for i in range(len(args)):
        offset = math.prod([len(a) for a in args[i:]]) // len(args[i])
        index = idx // offset
        elements += (args[i][index],)
        idx -= index * offset
    return elements

test


import itertools

def test_nth_product():
    a, b, c, d = [1, 2, 3, 4, 5, 6], [1, 2, 3, 4, 5], [1, 2, 3, 4], [1, 2, 3, 4, 5, 6, 7]
    for i, comb in enumerate(itertools.product(*[a, b, c, d])):
        assert comb == nth_product(i, a, b, c, d)
1 Like

Hi. more_itertools has an implementation.

1 Like

I know the implementation from more-itertools and it is based on islice which iterates through every combination until the index enumerator matches. As a common use case for nth_product is to create a “comb” through a very large solution landscape avoiding iteration is often a requirement. This can be for starting the population for a genetic algorithm or assuring an broad scan of the solution landscape for a Monte-Carlo simulation. I wrote this :arrow_up: because it was too slow.

Why don’t you try to contribute your implementation to more_itertools? This is how open source software improves.

3 Likes

Good idea. Let’s close this thread.