想要将itertools的输出存储为数组

Question

我有以下代码

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

for each in result:
    print(each)

带输出：

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

我想将“结果”中的各个“项目”存储为numpy数组，如果它们总结为2，我不确定itertools输出的确切数据类型。

伪代码中的示例：

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

Answer 1

所以使用理解：

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

Answer 2

这是一个功能性解决方案：

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