python - 向量化 numpy 索引并应用函数构建矩阵

Question

我有一个大小为 (d,N) 的矩阵 X。换句话说，有 N 个向量，每个向量有 d 个维度。例如，

X = [[1,2,3,4],[5,6,7,8]]

有 N=4 个 d=2 维的向量。

此外，我还有 rag 数组(列表列表)。索引是 X 矩阵中的索引列。例如，

I = [ [0,1], [1,2,3] ]

I[0]=[0,1] 索引矩阵 X 中的第 0 列和第 1 列。类似地，元素 I[1] 索引第 1,2 列和第 3 列。请注意，I 的元素是不属于一样长!

我想做的是使用 I 中的每个元素对矩阵 X 中的列进行索引，对向量求和并得到一个向量。对 I 的每个元素重复此操作，从而构建一个新矩阵 Y。矩阵 Y 应具有与 I 数组中的元素一样多的 d 维向量。在我的示例中，Y 矩阵将具有 2 个二维向量。

在我的示例中，元素 I[0] 指示从矩阵 X 中获取第 0 列和第 1 列。将矩阵 X 的两个二维向量相加并将该向量放入 Y(第 0 列)中。然后，元素 I[1] 告诉我们对矩阵 X 的第 1、2 和 3 列求和，并将这个新向量放入 Y(第 1 列)中。

我可以使用循环轻松完成此操作，但如果可能的话，我想将此操作向量化。我的矩阵 X 有数十万列，索引矩阵 I 有数万个元素(每个元素都是一个简短的索引列表)。

我的循环代码:

Y = np.zeros( (d,len(I)) )
for i,idx in enumerate(I):
    Y[:,i] = np.sum( X[:,idx], axis=1 )

Answer 1

这是一种方法-

# Get a flattened version of indices
idx0 = np.concatenate(I)

# Get indices at which we need to do "intervaled-summation" along axis=1
cut_idx = np.append(0,map(len,I))[:-1].cumsum()

# Finally index into cols of array with flattend indices & perform summation
out = np.add.reduceat(X[:,idx0], cut_idx,axis=1)

逐步运行-

In [67]: X
Out[67]: 
array([[ 1,  2,  3,  4],
       [15,  6, 17,  8]])

In [68]: I
Out[68]: array([[0, 2, 3, 1], [2, 3, 1], [2, 3]], dtype=object)

In [69]: idx0 = np.concatenate(I)

In [70]: idx0 # Flattened indices
Out[70]: array([0, 2, 3, 1, 2, 3, 1, 2, 3])

In [71]: cut_idx = np.append(0,map(len,I))[:-1].cumsum()

In [72]: cut_idx # We need to do addition in intervals limited by these indices
Out[72]: array([0, 4, 7])

In [74]: X[:,idx0]  # Select all of the indexed columns
Out[74]: 
array([[ 1,  3,  4,  2,  3,  4,  2,  3,  4],
       [15, 17,  8,  6, 17,  8,  6, 17,  8]])

In [75]: np.add.reduceat(X[:,idx0], cut_idx,axis=1)
Out[75]: 
array([[10,  9,  7],
       [46, 31, 25]])