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

时间:2019-10-21 23:45:21

标签: python numpy numpy-ndarray

我有一个形状为(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]]

我只想选择第二列的值为59的行。

为此,我创建了一个列表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 innp.wherenp.delete结合使用,这样我的问题就解决了!

1 个答案:

答案 0 :(得分:0)

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

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)