循环遍历多维数组并根据条件过滤答案

【问题标题】：Looping through multi-dimensional array and filtering based on a condition循环遍历多维数组并根据条件过滤
【发布时间】：2019-10-03 12:13:14
【问题描述】：

我相信有人提出了类似的问题，但没有一个涉及我所面临的问题。

我有一个形状为(H,W,L) 的数组 - 我必须遍历数组的每个实例以过滤出满足特定条件的值（它们的 x、y 位置）。（比如val > t_r 和val < t_c） - 我必须对每个 K 值重复此操作。

例如：如果我们有一个形状为 (2,3,4) 的数组

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

假设第一个条件是val > 2，第二个条件是val < 6，并将(row, col) 值存储在N x 3 数组中。其中前两个值是'row','col'，最后一个对应于层/第三维。那么操作的预期输出应该类似于 -

output = [[0,2,0],[1,0,0],[1,0,0],[1,1,0],[1,2,0]....] 这将对应于从A[:,:,0] 过滤的值

我想到的一种方法是 - 使用 3 个 for 循环 - i,j,k 循环遍历每个元素，但我无法弄清楚确切的实现。我还想尽可能实现矢量化。我可以使用一些指导。

【问题讨论】：

只是一般性评论：您几乎不需要 NumPy 循环。

标签： python arrays python-3.x numpy

【解决方案1】：

您可以使用np.nonzero 并矢量化您的比较。

a = np.asarray(A)
res = np.vstack(np.nonzero((a>2)&(a<6))).T

array([[0, 0, 2],
       [0, 1, 0],
       [0, 1, 1],
       [0, 1, 2],
       [1, 1, 1],
       [1, 1, 2],
       [2, 0, 0],
       [2, 1, 1],
       [2, 1, 2],
       [3, 0, 2],
       [3, 1, 0]], dtype=int64)

您可以随时根据自己的喜好重新排列列，例如：

res[:, [1,2,0]]

array([[0, 2, 0],
       [1, 0, 0],
       [1, 1, 0],
       [1, 2, 0],
       [1, 1, 1],
       [1, 2, 1],
       [0, 0, 2],
       [1, 1, 2],
       [1, 2, 2],
       [0, 2, 3],
       [1, 0, 3]], dtype=int64)

【讨论】：