Как выбрать строки в ndarray на основе условия списка - PullRequest
Новая
       67

Как выбрать строки в ndarray на основе условия списка

0 голосов
/ 22 октября 2019

У меня есть ndarray формы (1, 1, 5, 7), и это выглядит так: 

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, поэтому моя проблема решена сейчас!

1 Ответ

0 голосов
/ 22 октября 2019

Сначала вы должны изменить массив на 2-мерный, попробуйте это: 

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)
...