Search code examples
pythonpython-itertools

Want to store output from itertools as array

I have the following code

result = itertools.combinations_with_replacement(range(3),3)

for each in result:
    print(each)

with output:

(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 1, 1)
(0, 1, 2)
(0, 2, 2)
(1, 1, 1)
(1, 1, 2)
(1, 2, 2)
(2, 2, 2)

I want to store the individual "items" in "result" as an numpy arrays if they sum up to lets say 2. I'm not sure exactly what data type itertools outputs.

Example in pseudo code:

for each in result:
    if sum(each)==2:
        numpy array = each
Solution

  • So use the comprehension:

    import itertools
import numpy as np

result = itertools.combinations_with_replacement(range(3),3)
desired = [np.array(i) for i in result if sum(i)==2]
desired
#[array([0, 0, 2]), array([0, 1, 1])]