Shahroozevsky
Shahroozevsky

Reputation: 343

Element at index in Itertools.product

Can we get the element in a specific index from itertools.product result in Python 3? Like Below:

xlist = itertools.product('abc', repeat=3)
element = xlist[10]

UPDATE
The whole question turned upsidedown! I found that generating all set and look for an index is a big mistake! I looked at the possible duplicate of my question, but I don't get the answer!

Upvotes: 5

Views: 2471

Answers (4)

javidcf
javidcf

Reputation: 59711

The other solutions produce the correct result, but if you just want one element of the Cartesian product, iterating the generator returned by itertools.product is not the most efficient solution. You can directly build the item that you need without having to go through all the elements with a function like this:

from collections.abc import Sequence

def product_item(idx, *seqs, repeat=None):
    # Ensure inputs are actual sequences (list, tuple, str...)
    seqs = [seq if isinstance(seq, Sequence) else list(seq) for seq in seqs]
    # Repeat if needed
    if repeat is not None:
        seqs = seqs * repeat
    # Compute how many items does it take to advance on each sequence
    step = 1
    for seq in seqs:
        step *= len(seq)
    # Build product item
    item = [None] * len(seqs)
    for i, seq in enumerate(seqs):
        step //= len(seq)
        seq_idx = idx // step
        idx %= step
        item[i] = seq[seq_idx]
    return tuple(item)

print(product_item(10, 'abc', repeat=3))
# ('b', 'a', 'b')

The complexity for this solution is O(1). A quick comparison:

import itertools

# Solution with islice
product_item_islice = lambda idx, *seqs, repeat=None: next(
    itertools.islice(itertools.product(*seqs, repeat=repeat), idx, None))

idx = 100_000_000
seqs = ['abcdefgh']
repeat = 10
print(product_item(idx, *seqs, repeat=repeat))
# ('a', 'f', 'h', 'f', 'd', 'g', 'a', 'e', 'a', 'a')
print(product_item_islice(idx, *seqs, repeat=repeat))
# ('a', 'f', 'h', 'f', 'd', 'g', 'a', 'e', 'a', 'a')

%timeit product_item(idx, *seqs, repeat=repeat)
# 3.7 µs ± 46.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit product_item_islice(idx, *seqs, repeat=repeat)
# 448 ms ± 7.55 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Upvotes: 3

tobias_k
tobias_k

Reputation: 82919

You do not have to iterate the generator until you get to the index, but you can generate that product immediately, in O(1) as described here. Adapting that to the (lst, repeat) form:

def nth_product(lst, repeat, n):
    k = len(lst)
    return [lst[n // (k**r) % k] for r in range(repeat-1, -1, -1)]

Example:

>>> lst = list(range(10))
>>> ref = list(itertools.product(lst, repeat=3))
>>> all(nth_product(lst, 3, i) == list(r) for i, r in enumerate(ref))
True

Upvotes: 3

Gareth Latty
Gareth Latty

Reputation: 89027

The result of itertools.product() isn't a list, but rather an iterable - which is why this doesn't work, you can't access an iterable by index.

You can make a list using list() - but this means computing all the values, which could be very inefficient. In the example case, while we have to compute all the values leading up to the one we want, we don't need to store them all in memory, nor do we need to compute the rest after we get the one we want.

If you only want one element, the better solution is to only consume the part you need - this can be done easily with itertools.islice().

element = next(itertools.islice(xlist, 10, None))

As the slice is another iterable, we use next() to get the first item.

islice functions very similarly to the list slice (as the name implies). Note that once you have consumed some of the iterable, working with it further will work from where you left off. You should either get what you need in a single iteration, or create a list (or other data structure if appropriate).

The big advantage of islice over other ways of throwing away the initial values you are not interested in is that it is implemented very efficiently within Python, and so is likely to be the fastest option, as well as being flexible if you later need more than just the one element.

Upvotes: 2

Dani Mesejo
Dani Mesejo

Reputation: 61910

Assuming you refer to Python 3, you could use enumerate to filter to index the generator and then use next:

import itertools

it = enumerate(itertools.product('abc', repeat=3))
result = next(e for i, e in it if i == 10)
print(result)

Output

('b', 'a', 'b')

Upvotes: 2

Related Questions