我需要扩展this question，它根据第二个数组的索引对数组的值求和。让A为结果数组，B为索引数组，C为要求的数组。然后A[i] = sum超过C，index(B) == i。
相反，我的设置是
N = 5
M = 2
A = np.zeros((M,N))
B = np.random.randint(M, size=N) # contains indices for A
C = np.random.rand(N,N)
我需要A[i,j] = sum_{k in 0...N} C[j,k]使C[k] == i，即以B匹配i的索引为条件的rowum。有没有一种有效的方法来做到这一点？对于我的应用程序，N大约为10,000，M大约为20.在最小化问题中每次迭代都会调用此操作...我当前的循环方法非常慢。
谢谢！
答案 0 :(得分：4)
根据@ DSM的评论，我假设你的C[k] == i应该是B[k] == i。如果是这种情况，您的循环版本是否看起来像这样？
import numpy as np
N = 5
M = 2
A = np.zeros((M,N))
B = np.random.randint(M, size=N) # contains indices for A
C = np.random.rand(N,N)
for i in range(M):
for j in range(N):
for k in range(N):
if B[k] == i:
A[i,j] += C[j,k]
有多种方法可以将此问题矢量化。我将在下面展示我的思考过程，但是有更有效的方法可以做到这一点（例如，@ DSM的版本可以识别问题中固有的矩阵乘法）。
为了便于解释，这里是一种方法的演练。
让我们从重写内部k循环开始：
for i in range(M):
for j in range(N):
A[i,j] = C[j, B == i].sum()
可能更容易将其视为C[j][B == i].sum()。我们只是选择j的{{1}} th 行，只选择C等于B的那一行中的元素，然后求和它们。
接下来让我们分解外部i循环。不幸的是，现在我们将要达到可读性开始受损的程度......
i这里有几个不同的技巧。在这种情况下，我们将迭代i = np.arange(M)[:,np.newaxis]
mask = (B == i).astype(int)
for j in range(N):
A[:,j] = (C[j] * mask).sum(axis=-1)
的列。 A的每列是A的相应行的子集的总和。 C行的子集由C等于行索引B的位置确定。
为了绕过i，我们通过向i添加新轴来制作一个二维数组i。 （如果您对此感到困惑，请查看numpy broadcasting的文档。）换句话说：
B == i我们想要的是获取B:
array([1, 1, 1, 1, 0])
i:
array([[0],
[1]])
B == i:
array([[False, False, False, False, True],
[ True, True, True, True, False]], dtype=bool)
的两个（M）过滤总和，C[j]中每行一个。这将为我们提供一个与B == i中的j th 列对应的双元素向量。
我们不能通过直接索引A来做到这一点，因为结果不会保持它的形状，因为每行可能有不同数量的元素。我们会通过将C掩码乘以B == i的当前行来解决此问题，从而导致C为B == i的零，以及当前行中的值False这是真的。
为此，我们需要将布尔数组C转换为整数：
B == i所以当我们将它乘以当前行mask = (B == i).astype(int):
array([[0, 0, 0, 0, 1],
[1, 1, 1, 1, 0]])
：
C然后我们可以对每一行进行求和以得到C[j]:
array([ 0.19844887, 0.44858679, 0.35370919, 0.84074259, 0.74513377])
C[j] * mask:
array([[ 0. , 0. , 0. , 0. , 0.74513377],
[ 0.19844887, 0.44858679, 0.35370919, 0.84074259, 0. ]])
的当前列（当它被分配给A时，它将被广播到列中：
A[:,j]最后，打破最后一个循环，我们可以应用完全相同的原则为(C[j] * mask).sum(axis=-1):
array([ 0.74513377, 1.84148744])
上的循环添加第三个维度：
j正如@DSM建议的那样，您也可以这样做：
i = np.arange(M)[:,np.newaxis,np.newaxis]
mask = (B == i).astype(int)
A = (C * mask).sum(axis=-1)
对于大多数A = (B == np.arange(M)[:,np.newaxis]).dot(C.T)
和M来说，这是迄今为止最快的解决方案，可以说是最优雅的（比我的解决方案更优雅）。
让我们分解一下。
N完全等同于上面“向量化最外圈”部分中的B == np.arange(M)[:,np.newaxis]。
关键是要认识到所有B == i和j循环都等同于矩阵乘法。 k会将布尔dot数组转换为与幕后B == i相同的dtype，因此我们无需担心将其明确地转换为其他类型。
之后，我们只是在C（一个5x5数组）的转置和上面的“mask”0和1数组上执行矩阵乘法，产生一个2x5数组。
C将利用您已安装的任何优化BLAS库（例如dot，ATLAS），因此它非常快。
对于小MKL和M，差异不太明显（循环和DSM版本之间约为6倍）：
N然而，一旦M, N = 2, 5
%timeit loops(B,C,M)
10000 loops, best of 3: 83 us per loop
%timeit k_vectorized(B,C,M)
10000 loops, best of 3: 106 us per loop
%timeit vectorized(B,C,M)
10000 loops, best of 3: 23.7 us per loop
%timeit askewchan(B,C,M)
10000 loops, best of 3: 42.7 us per loop
%timeit einsum(B,C,M)
100000 loops, best of 3: 15.2 us per loop
%timeit dsm(B,C,M)
100000 loops, best of 3: 13.9 us per loop
和M开始增长，差异就变得非常显着（~600x）（注意单位！）：
N答案 1 :(得分：3)
我假设@DSM发现了你的拼写错误，你想要：
A[i,j] = sum_{k in 0...N} C[j,k] where B[k] == i
然后，您可以循环i in range(M)，因为M相对较小。
A = np.array([C[:,B == i].sum(axis=1) for i in range(M)])