如何打印以特定数字开头的数组元素？

Question

我有一个Numpy数组

a= [[[ 6, -1, -1]],

       [[ 8, -1,  3]],

       [[10,  1,  1]],

       [[10, -1, -1]],

       [[10,  1, -1]]]

我希望只保留那些以10开头的行;所以我希望我的输出数组是

a_out= [[[10,  1,  1]],

           [[10, -1, -1]],

           [[10,  1, -1]]]

我正在尝试迭代数组的第一个元素，但是如果我尝试a[0,0:]或a[0:,0]，我会得到整个行或列。

Answer 1

使用布尔掩码：

>>> arr
array([[[ 6, -1, -1]],

       [[ 8, -1,  3]],

       [[10,  1,  1]],

       [[10, -1, -1]],

       [[10,  1, -1]]])
>>> mask = arr[:,:,0] == 10
>>> arr[mask]
array([[[10,  1,  1]],

       [[10, -1, -1]],

       [[10,  1, -1]]])

Answer 2

确保您的a是一个numpy数组而不是Python列表。然后你可以使用它：

a[a[:,:,0] == 10]

其中内部是bool列表以选择行。

Answer 3

如果是列表而不是数组，则可以执行以下操作：

filter(lambda x: x[0][0] == 10, a)