时间戳

首页 > 解决方案 > 3D 数组中最后两个非零元素的平均值

python - 3D 数组中最后两个非零元素的平均值

问题描述

我有一个 (n by i by j) - 3D numpy 数组：a_3d_array（2 x 5 x 3）

array([[[1, 2, 3],
    [1, 1, 1],
    [2, 2, 2],
    [0, 3, 3],
    [0, 0, 4]],

   [[1, 2, 3],
    [2, 2, 2],
    [3, 3, 3],
    [0, 4, 4],
    [0, 0, 5]]]).

对于 n 中的每一列 j，我想提取最后 2 个非零元素并计算平均值，然后将结果放入 (n by j) 数组中。我目前所做的是使用 for 循环

import numpy as np

a_3d_array = np.array([[[1, 2, 3],
                        [1, 1, 1],
                        [2, 2, 2],
                        [0, 3, 3],
                        [0, 0, 4]],
                       [[1, 2, 3],
                        [2, 2, 2],
                        [3, 3, 3],
                        [0, 4, 4],
                        [0, 0, 5]]])

aveCol = np.zeros([2,3])
for n in range(2):
    for j in range(3):
        temp = a_3d_array[n,:,j]
        nonzero_array = temp[np.nonzero(temp)]
        aveCol[n, j] = np.mean(nonzero_array[-2:])

得到想要的结果

print(aveCol)
[[1.5 2.5 3.5] [2.5 3.5 4.5]]

效果很好。但我想知道是否有更好的 Pythonic 方式来做同样的事情？

我发现与我的问题最相似的是这里。但我不太明白在稍微不同的背景下解释的答案。

标签： pythonnumpymultidimensional-arraymean

解决方案

TL;DR 据我所知，安的回答是最快的

每个m都是一个n×i 2D 数组，接下来我们r对其转置，即执行计算的“列” - 在这个“列”上，我们丢弃所有零，我们将最后两个非零元素相加取平均值

In [17]: np.array([[sum(r[r!=0][-2:])/2 for r in m.T] for m in a])
Out[17]: 
array([[1.5, 2.5, 3.5],
       [2.5, 3.5, 4.5]])

编辑1

看起来它比你的循环更快

In [19]: %%timeit
    ...: avg = np.zeros([2,3])
    ...: for n in range(2):
    ...:     for j in range(3):
    ...:         temp = a[n,:,j]
    ...:         nz = temp[np.nonzero(temp)]
    ...:         avg[n, j] = np.mean(nz[-2:])
95.1 µs ± 596 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [20]: %timeit np.array([[sum(r[r!=0][-2:])/2 for r in m.T] for m in a])
45.5 µs ± 394 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

编辑2

In [22]: %timeit np.array([[np.mean(list(filter(None, a[n,:,j]))[-2:]) for j in range(3)] for n in range(2)])
145 µs ± 689 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

编辑3

In [25]: %%timeit
    ...: i = np.indices(a.shape)
    ...: i[:, a == 0] = -1
    ...: i = np.sort(i, axis=2)
    ...: i = i[:, :, -2:, :]
    ...: a[tuple(i)].mean(axis=1)
64 µs ± 239 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Edit4 突发新闻信息

安回答的罪魁祸首是np.mean！！

In [29]: %timeit np.array([[sum(list(filter(None, a[n,:,j]))[-2:])/2 for j in range(3)] for n in range(2)])
32.7 µs ± 111 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

推荐阅读