沿 i 个元素求和 3d numpy 数组 [i,j]

Question

对于形状为 ([N,N,3]) 的 numpy 数组，

[[1,2,3], [1,2,3],[1,2,3]],
[[4,5,6],[4,5,6],[4,5,6]],
[[7,8,9],[7,8,9],[7,8,9]]

我怎样才能得到一个 [N,3] 与 j 元素总结看起来像这样？

[2,4,6], [8,10,12], [14,16,18]

我正在考虑重塑它，尽管我认为没有必要。 我知道我想垂直添加，所以总和轴 = 0。 我只是发现很难保持结构不变。

Answer 1

您将乘以 2，但对于总和，这是一个解决方案：

# Convert to list
>>> lst = a.tolist()

# Sum by index
>>> sum_by_index = lambda x: [x[0][i] + x[1][i] + x[2][i] for i in range(3)]

# For multiplication
>>> mply_by_index = lambda x: [x[0][i] * 2 for i in range(3)]

# Apply
>>> list(map(sum_by_index, lst))
[[3.0, 6.0, 9.0], ...]

>>> list(map(mply_by_index, lst))
[[2.0, 4.0, 6.0], ...]

Answer 2

我怎样才能得到一个 [N,3] 与 j 元素总结看起来像这样？

这是你想要的？

j = 2

a[:,:j,:].sum(1)

array([[ 2,  4,  6],
       [ 8, 10, 12],
       [14, 16, 18]])