python - Numpy 根据列表折叠列

Question

在 NumPy 中，我有一个 d x n 数组 A 和一个长度为 n 的列表 L ，描述我希望 A 的每一列在矩阵 B 中结束的位置。这个想法是矩阵 B 的列 i 是 A 的所有列的 sum 其相应的值在 L 中是 i。

我可以用 for 循环来做到这一点:

A = np.arange(15).reshape(3,5)
L = [0,1,2,1,1]
n_cols = 3
B = np.zeros((len(A), n_cols)) 
# assume I know the desired number of columns, 
# which is also one more than the maximum value of L
for i, a in enumerate(A.T):
    B[:, L[i]] += a

我想知道是否有办法通过切片数组 A(或以其他矢量化方式)来做到这一点？

Answer 1

您正在使用 L 对 A 中的列进行求和/折叠以选择这些列。此外，您正在根据 L 元素的唯一性更新输出数组的列。

因此，您可以使用 np.add.reduceat对于矢量化解决方案，就像这样 -

sidx = L.argsort()
col_idx, grp_start_idx = np.unique(L[sidx],return_index=True)
B_out = np.zeros((len(A), n_cols))
B_out[:,col_idx] = np.add.reduceat(A[:,sidx],grp_start_idx,axis=1)

---

运行时测试-

In [129]: def org_app(A,n_cols):
     ...:     B = np.zeros((len(A), n_cols)) 
     ...:     for i, a in enumerate(A.T):
     ...:         B[:, L[i]] += a
     ...:     return B
     ...: 
     ...: def vectorized_app(A,n_cols):
     ...:     sidx = L.argsort()
     ...:     col_idx, grp_start_idx = np.unique(L[sidx],return_index=True)
     ...:     B_out = np.zeros((len(A), n_cols))
     ...:     B_out[:,col_idx] = np.add.reduceat(A[:,sidx],grp_start_idx,axis=1)
     ...:     return B_out
     ...: 

In [130]: # Setup inputs with an appreciable no. of cols & lesser rows
     ...: # so as that memory bandwidth to work with huge number of 
     ...: # row elems doesn't become the bottleneck
     ...: d,n_cols = 10,5000
     ...: A = np.random.rand(d,n_cols)
     ...: L = np.random.randint(0,n_cols,(n_cols,))
     ...: 

In [131]: np.allclose(org_app(A,n_cols),vectorized_app(A,n_cols))
Out[131]: True

In [132]: %timeit org_app(A,n_cols)
10 loops, best of 3: 33.3 ms per loop

In [133]: %timeit vectorized_app(A,n_cols)
100 loops, best of 3: 1.87 ms per loop

随着行数变得与 A 中的列数相当，矢量化方法的高内存带宽要求将抵消它的任何明显加速。