gpt4 book ai didi

python - numpy 中 itertools.combinations 的 N-D 版本

转载 作者:太空狗 更新时间:2023-10-29 16:55:06 27 4
gpt4 key购买 nike

我想实现 itertools.combinations对于 NumPy 。基于this discussion ,我有一个适用于一维输入的函数:

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],并且 combs([4, 5, 6], 2) 生成 [4, 5], [4, 6], [5, 6],然后是 combs((1,2,3) and (4,5,6), 2) 应该生成 [1, 2], [1, 3], [2, 3] 和 [4, 5], [4, 6], [5, 6] 其中“and”仅表示平行的行或列(以有意义的为准)。 (同样适用于其他维度)

我不确定:

  1. 如何使维度以与其他函数的工作方式一致的逻辑方式工作(比如一些 numpy 函数如何具有 axis= 参数,以及轴 0 的默认值。所以可能是 axis 0应该是我合并的那一个,其他所有轴都代表并行计算?)
  2. 如何让上面的代码与 ND 一起工作(现在我得到 ValueError: setting an array element with a sequence.)
  3. 有没有更好的方法来做 dt = dtype([('', a.dtype)]*r)

最佳答案

您可以使用 itertools.combinations() 创建索引数组,然后使用 NumPy 的花式索引:

import numpy as np
from itertools import combinations, chain
from scipy.special import comb

def comb_index(n, k):
count = comb(n, k, exact=True)
index = np.fromiter(chain.from_iterable(combinations(range(n), k)),
int, count=count*k)
return index.reshape(-1, k)

data = np.array([[1,2,3,4,5],[10,11,12,13,14]])

idx = comb_index(5, 3)
print(data[:, idx])

输出:

[[[ 1  2  3]
[ 1 2 4]
[ 1 2 5]
[ 1 3 4]
[ 1 3 5]
[ 1 4 5]
[ 2 3 4]
[ 2 3 5]
[ 2 4 5]
[ 3 4 5]]

[[10 11 12]
[10 11 13]
[10 11 14]
[10 12 13]
[10 12 14]
[10 13 14]
[11 12 13]
[11 12 14]
[11 13 14]
[12 13 14]]]

关于python - numpy 中 itertools.combinations 的 N-D 版本,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/16003217/

27 4 0
Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
广告合作:1813099741@qq.com 6ren.com