我想实施itertools.combinations对于 numpy.基于这次讨论,我有一个适用于一维输入的函数:
def combs(a, r):
"""
Return successive r-length combinations of elements in the array a.
Should produce the same output as array(list(combinations(a, r))), but
faster.
"""
a = asarray(a)
dt = dtype([('', a.dtype)]*r)
b = fromiter(combinations(a, r), dt)
return b.view(a.dtype).reshape(-1, r)
输出是有意义的:
In [1]: list(combinations([1,2,3], 2))
Out[1]: [(1, 2), (1, 3), (2, 3)]
In [2]: array(list(combinations([1,2,3], 2)))
Out[2]:
array([[1, 2],
[1, 3],
[2, 3]])
In [3]: combs([1,2,3], 2)
Out[3]:
array([[1, 2],
[1, 3],
[2, 3]])
然而,如果我可以将其扩展到 N 维输入,那就更好了,其中额外的维度可以让您快速地一次执行多个调用。所以,从概念上讲,如果combs([1, 2, 3], 2)
产生[1, 2], [1, 3], [2, 3]
, and combs([4, 5, 6], 2)
产生[4, 5], [4, 6], [5, 6]
, then combs((1,2,3) and (4,5,6), 2)
应该产生[1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6]
其中“and”仅代表平行的行或列(以有意义的为准)。 (对于其他尺寸同样如此)
我不知道:
- 如何使维度以与其他函数的工作方式一致的逻辑方式工作(例如某些 numpy 函数如何具有
axis=
参数,默认为轴 0。所以轴 0 可能应该是我要组合的轴,而所有其他轴仅代表并行计算?)
- 如何让上面的代码与 ND 一起工作(现在我明白了
ValueError: setting an array element with a sequence.
)
- 有更好的方法吗
dt = dtype([('', a.dtype)]*r)
?