numpy 按索引求和矩阵行答案

【问题标题】：numpy summing matrix-rows by indicesnumpy 按索引求和矩阵行
【发布时间】：2018-10-13 22:15:53
【问题描述】：

我有 3 个矩阵（np 数组）：
A 的形状为 (n,m) ； B 的形状为 (m,k)； C 的形状为 (n,k)

矩阵 C 仅具有来自集合 {-1,0,1} 的值，它是某种“指标”：如果 C[i,j]==1 那么我想添加第 i 行A到b的第j列；如果 C[i,j]==(-1) 然后减去（0 什么都不做）。

它可以很容易地用循环来完成，但我想知道是否有一种可能更快的矢量化方式来完成它？

示例代码：

C = np.array([[-1,  0,  0,  0,  1],
              [ 0,  0,  0,  0, -1],
              [ 0,  0,  0,  0, -1],
              [-1,  0,  0,  1,  1]])
a,b = np.where(C==1)
#here a=[0,3,3] and b=[4,3,4]
A[a,:] = [[0, 1, 2, 3, 4, 5, 6],
          [3, 3, 3, 3, 3, 3, 3],
          [3, 3, 3, 3, 3, 3, 3]]
B[:,b] += A[a]  #B is all 0.0 before

预期结果：

array([[ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  4.],
       [ 0.,  0.,  0.,  3.,  5.],
       [ 0.,  0.,  0.,  3.,  6.],
       [ 0.,  0.,  0.,  3.,  7.],
       [ 0.,  0.,  0.,  3.,  8.],
       [ 0.,  0.,  0.,  3.,  9.]])

实际结果：

array([[ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.]])

【问题讨论】：

之前的A是什么？
还有，B之前是什么？

标签： python numpy vectorization

【解决方案1】：

我们可以在B的转置视图上简单地使用np.add.at -

np.add.at(B.T, b, A[a])

示例运行 -

In [39]: C = np.array([[-1,  0,  0,  0,  1],
    ...:               [ 0,  0,  0,  0, -1],
    ...:               [ 0,  0,  0,  0, -1],
    ...:               [-1,  0,  0,  1,  1]])
    ...: a,b = np.where(C==1)
    ...: A = np.zeros((4,7),dtype=int)
    ...: A[a,:] = np.array([[0, 1, 2, 3, 4, 5, 6],
    ...:           [3, 3, 3, 3, 3, 3, 3],
    ...:           [3, 3, 3, 3, 3, 3, 3]])

In [40]: # Initialize B
    ...: B = np.zeros((7,5),dtype=int)

In [41]: np.add.at(B.T, b, A[a])

In [42]: B
Out[42]: 
array([[0, 0, 0, 3, 3],
       [0, 0, 0, 3, 4],
       [0, 0, 0, 3, 5],
       [0, 0, 0, 3, 6],
       [0, 0, 0, 3, 7],
       [0, 0, 0, 3, 8],
       [0, 0, 0, 3, 9]])

作为commented by @DSM，我们还可以使用矩阵乘法，从而避免获取C==1 索引的步骤-

A.T.dot(C==1)

【讨论】：

但是如果我们还必须处理 C[i,j]==(-1) 的情况，不就是A.T @ C吗？
@DSM 好吧，我们没有得到完整的A，所以我只使用给定的A[a]。但是，是的，我想我们可以在将所有 -1s 设置为 0 后做到这一点。
A 的行数应该与 C 相同。
@Kefeng91 如问题所述：A is of shape (n,m) ; and C is of shape (n,k)。
C 的形状为 (4, 5)，你的 A 的形状为 (10, 7)。

【解决方案2】：

我很难把这一切都记在心里，但我相信这就是你想要做的：

import numpy as np

n, m, k = 2, 3, 4 

# A is full of 1s
A = np.ones((n,m))
# B is full of 0s
B = np.zeros((m,k))
# C is...  that
C = np.array([
    [ 0,  0, -1,  1],
    [ 0, -1,  1,  0]
])

# Where to add A to B
add_x, add_y = np.where(C==1)
# Where to subtract A to B
sub_x, sub_y = np.where(C==-1)

# Select columns of B and += rows of A 
# (using .T to transpose A so it's the same shape)
B[:,add_y] += A[add_x,:].T
# ditto thing above, but for -=
B[:,sub_y] -= A[sub_x,:].T
print(B)

输出：

[[ 0. -1.  0.  1.]
 [ 0. -1.  0.  1.]
 [ 0. -1.  0.  1.]]

【讨论】：