NumPy：用三维数组中各元素的平均值替换第三维度上的所有元素

Question

我有一个三维数组，其尺寸为400*800*3。我想用该维度的平均值替换第三个维度中的所有元素。目前，我已经通过使用循环完成了这个任务。

test_data=np.random.randint(0,256,size=(400,800,3))
for i in range(400):
    for j in range(800):
        mn = np.mean(test_data[i,j])
        test_data[i,j]=mn

我想知道是否有一种高效且简洁的方法来实现相同的目标。谢谢。

另外，如果我想要三个元素的加权平均值而不是平均值，该怎么办？也就是说，最后一行应该被替换为：

test_data[i,j,0]=test_data[i,j,0]*0.2
test_data[i,j,1]=test_data[i,j,1]*0.5
test_data[i,j,2]=test_data[i,j,2]*0.3

- I Khalil

Answer 1

获取沿最后一个轴的平均值，并将其分配到所有位置 [:] 中 -

test_data[:] = test_data.mean(axis=-1,keepdims=1)

或者，我们可以创建一个新的数组，并在最后一个轴上进行复制 -

mean_vals = test_data.mean(axis=-1,keepdims=1).astype(test_data.dtype)
test_data_out = np.repeat(mean_vals,3,axis=-1)

如果需要一个只读版本并且速度更快的方式，请使用np.broadcast_to -

test_data_out = np.broadcast_to(mean_vals, test_data.shape) 

对于加权平均部分，如果您是指：

test_data[i,j]= test_data[i,j,0]*0.2 + test_data[i,j,1]*0.5 + test_data[i,j,2]*0.3

为此，我们可以使用多维张量缩减：np.tensordot。

W = [0.2,0.5,0.3] # weights
W_mean = np.tensordot(test_data,W, axes=((-1,-1)))[...,None]
test_data[:] = W_mean.astype(test_data.dtype)

相反，如果你的意思是：

test_data[i,j,0]=test_data[i,j,0]*0.2
test_data[i,j,1]=test_data[i,j,1]*0.5
test_data[i,j,2]=test_data[i,j,2]*0.3

因此，没有总和缩减，所以我们可以简单地利用广播 -

test_data[:] = (test_data*[0.2, 0.5, 0.3]).astype(test_data.dtype)

- Divakar

Answer 2

这应该可以工作：

test_data[:,:,0] += test_data[:,:,1]
test_data[:,:,0] += test_data[:,:,2]
test_data[:,:,0] *= 1.0/3
test_data[:,:,1] = test_data[:,:,0]*0.5
test_data[:,:,2] = test_data[:,:,0]*0.3
test_data[:,:,0] *= 0.2

这种方法可以“原地”工作，而不需要额外的内存。

- 6502

Answer 3

In [556]: test_data=np.arange(24.).reshape(2,4,3)
In [557]: test_data
Out[557]: 
array([[[  0.,   1.,   2.],
        [  3.,   4.,   5.],
        [  6.,   7.,   8.],
        [  9.,  10.,  11.]],

       [[ 12.,  13.,  14.],
        [ 15.,  16.,  17.],
        [ 18.,  19.,  20.],
        [ 21.,  22.,  23.]]])

沿着最后一个维度缩放值只是一种乘法操作：

In [558]: test_data*np.array([0.2, 0.5, 0.3])
Out[558]: 
array([[[  0. ,   0.5,   0.6],
        [  0.6,   2. ,   1.5],
        [  1.2,   3.5,   2.4],
        [  1.8,   5. ,   3.3]],

       [[  2.4,   6.5,   4.2],
        [  3. ,   8. ,   5.1],
        [  3.6,   9.5,   6. ],
        [  4.2,  11. ,   6.9]]])

可以使用 test_data *= np.array([....]) 直接实现。

对于最后一个维度取平均值只需：

In [559]: test_data.mean(axis=-1)
Out[559]: 
array([[  1.,   4.,   7.,  10.],
       [ 13.,  16.,  19.,  22.]])

这是2D。@divakar已经展示了如何保持它为3D（keep_dims）甚至是全尺寸（2,4,1）或者（2,4,3）。

In [561]: test_data.mean(axis=-1,keepdims=True)
Out[561]: 
array([[[  1.],
        [  4.],
        [  7.],
        [ 10.]],

       [[ 13.],
        [ 16.],
        [ 19.],
        [ 22.]]])

np.average 允许您指定权重（加权平均）：

In [568]: np.average(test_data, axis=-1, weights=[0.2,0.5,0.3])
Out[568]: 
array([[  1.1,   4.1,   7.1,  10.1],
       [ 13.1,  16.1,  19.1,  22.1]])

这与在最后一个维度上进行点积相同

In [569]: np.dot(test_data, [0.2,0.5,0.3])
Out[569]: 
array([[  1.1,   4.1,   7.1,  10.1],
       [ 13.1,  16.1,  19.1,  22.1]])

- hpaulj