首页 > 解决方案 > 如何根据列表条件选择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]]

我只想选择第二列的值为5or的行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我的问题现在解决了!

标签: pythonnumpynumpy-ndarray

解决方案


首先,您必须将数组更改为二维,试试这个:

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)

推荐阅读