如何根据列表条件在ndarray中选择行

Question

我有一个形状为（1，1，5，7）的ndarray，它看起来像这样：

array[0][0] = 
[[ 0.   5.   0.9  0.6  0.3  0.8  0.7]
 [ 0.   5.   0.5  0.2  0.2  0.4  0.7]
 [ 0.   9.   0.6  0.8  0.4  1.   1. ]
 [ 0.   9.   0.5  0.8  0.3  0.9  0.6]
 [ 0.  16.   0.6  0.1  0.   0.4  0.4]]

我只选择第二列的值为5或9的行。

为此，我创建了一个列表list = [5, 9]和下面的一个函数，以检查第二列的值是否为list中的值：

def check(list, array):
    tf_array = np.zeros(np.shape(array))
    for l in list:
        for i, element in enumerate(array):
            if element == l:
                tf_array[i] = True
    return tf_array

然后，我使用np.compress(check(list, array[0, 0, :, 1]), array, axis=2)来提取第二列的值是列表中任一列的行。

但是，我觉得必须有一个更简单的方法！抱歉，我刚接触Numpy，所以有人可以帮我吗？非常感谢！

编辑

感谢您的回答！抱歉，我忘记提及该列表是通用的，我希望能够指定它。我突然想起可以将条件not in与np.where和np.delete结合使用，所以我的问题现在解决了！

Answer 1

首先，您必须将数组更改为二维，请尝试以下操作：

import numpy as np

a=np.array([[ 0,   5,  0.9, 0.6,  0.3,  0.8,  0.7],
 [ 0,   5,   0.5, 0.2,  0.2,  0.4,  0.7],
 [ 0,   9,   0.6,  0.8,  0.4,  1,   1 ],
 [ 0,   9,   0.5,  0.8,  0.3,  0.9,  0.6],
 [ 0,  16,   0.6,  0.1,  0,   0.4,  0.4]])

for l in a:
    if l[1]==5 or l[1]==9:
        print(l)