Python中所有向量组合的集合

时间:2013-09-05 01:11:48

标签: python math numpy cartesian

我在创建可以被视为Python中的矢量数组的笛卡尔积的问题时遇到了问题。我有一个代码,它通过r变量给出n个数的所有可能分区,并将其作为numpy数组返回。我想做的是能够调用该代码任意次数,然后生成一组所有可能的数组组合。

举一个例子,我可以调用分区代码和每次连续调用(对于变化的参数集)

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

我正在寻找的是能够返回集合

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

作为整体数组,或逐行返回(由于明显的内存问题,因为分区数量的大小增加)。

以前我使用itertools.product解决了这个问题,并在PyPy下运行代码。但是,由于在项目的其他部分需要Numpy,我不得不从PyPy切换到标准python,并且我试图通过使用Numpy来复制PyPy代码的速度。我已经设法让这个工作非常粗略,但代码花了很多时间在数据类型之间进行更改,以便尝试将解决方案一起引导,这对于实现来说是不切实际的。

我想知道是否有人能够帮助我提供一些关于如何在Python中取得进展的指导。

由于

1 个答案:

答案 0 :(得分:1)

这应该让你开始:

import numpy as np
import itertools as it

def row_product(*arrays):
    lengths = np.array([x.shape[0] for x in arrays])
    positions = np.cumsum(lengths)

    ranges = np.arange(positions[-1])
    ranges = np.split(ranges,positions[:-1])

    total = np.concatenate((arrays),axis=0)

    inds = np.fromiter(it.chain.from_iterable(it.product(*ranges)), np.int)
    inds = inds.reshape(-1, len(arrays))

    return np.take(total, inds, axis=0)

最后一个维度必须相同。

显示结果:

a=np.array([[2,0],[1,1],[2,0]])
b=np.array([[1,0],[0,1]])
c=np.array([[0,0]])

print row_product(a,b,c)

[[[2 0]
  [1 0]
  [0 0]]

 [[2 0]
  [0 1]
  [0 0]]

 [[1 1]
  [1 0]
  [0 0]]

 [[1 1]
  [0 1]
  [0 0]]

 [[2 0]
  [1 0]
  [0 0]]

 [[2 0]
  [0 1]
  [0 0]]]

这是一个3D阵列,其中唯一的组合位于最后两个轴中。似乎相当快,1M独特的组合大约需要1/6秒。