CODEWITHSUNDEEP
pythonpython-itertools
kPuck
kPuck

Reputation: 147

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

Upvotes: 1

Views: 477

Answers (2)

zipa
zipa

Reputation: 27869

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])]

Upvotes: 2

jpp
jpp

Reputation: 164693

Here is a functional solution:

import itertools
import numpy as np

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

list(map(np.array, filter(lambda x: sum(x)==2, result)))

# [array([0, 0, 2]), array([0, 1, 1])]

Upvotes: 1

Related Questions