给定一个numpy数组,例如
A = np.ones(shape=(7, 6), dtype=np.float32)
和一个列表
v = [[0, 2], [1, 4], [3, 5, 6]]
我想做的是在给定v
中的每个项目的情况下,在A
中添加行,对于v
,有3个项目;对于v[0]
,明智地添加行0
和行2
列。输出的形状为(3, 6)
,输出为
res = array([[2., 2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2., 2.],
[3., 3., 3., 3., 3., 3.]])
# res[0] = A[0] + A[2]
# res[1] = A[1] + A[4]
# res[2] = A[3] + A[5] + A[6]
这里有一个更清楚的例子,给出一个矩阵
m = [[1, 2, 3],
[2, 3, 4],
[1, 1, 1],
[2, 2, 2],
[1, 1, 1]]
和要添加的行
v = [[0, 1, 3], [2]]
因此,这里将矩阵m
中的0
、1
和3
行相加,并且由于在[2]
中只有一行要相加,所以结果是
# res.shape = (2, 3)
res[0] = m[0] + m[1] + m[3]
res[1] = m[2]
还有什么更优雅的方法吗?
您可以使用花式索引从数组中选择行。
对于A
:
>>> A = np.ones(shape=(7, 6), dtype=np.float32)
>>> v = [[0, 2], [1, 4], [3, 5, 6]]
>>> np.array([A[rows].sum(axis=0) for rows in v])
array([[2., 2., 2., 2., 2., 2.],
[2., 2., 2., 2., 2., 2.],
[3., 3., 3., 3., 3., 3.]], dtype=float32)
对于m
:
>>> m = np.array([[1, 2, 3], [2, 3, 4], [1, 1, 1], [2, 2, 2], [1, 1, 1]])
>>> v = [[0, 1, 3], [2]]
>>> np.array([m[rows].sum(axis=0) for rows in v])
array([[5, 7, 9],
[1, 1, 1]])
我不知道这是否可以进一步向量化。