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)