Индексный список массивов, сохраняющих размерность в numpy

Question

У меня есть список массивов:

data = [array([4,2,3,4], dtype=uint16),
        array([6.6, 7.4, 5.0, 9.5], dtype=float32)] 

Я хотел бы получить значения из приведенного выше списка массивов, которые удовлетворяют условию, например:

condition = array([True, False, False, True])

Таким образом, что результат будет следующим:

data[:][condition]
 # Equals to -> [array([4,4], dtype=uint16), 
 #               array([6.6, 9.5], dtype=float32)]

При сохранении той же формы, очевидно, будет уменьшено количество значений

Я знаю, что делать:

data[0][np.where(condition)]

Это дает мне то, что я хочу, но только для этого массива [0].

Как я могу сделать это для нескольких таких массивов?

Answer 1

Если все ваши массивы в списке имеют одинаковую форму, самый элегантный способ - преобразовать ваш список в numpy массив и использовать numpy, индексируя как @Quang, упомянутое в комментариях:

data = [np.array([4,2,3,4], dtype=np.uint16),
        np.array([6.6, 7.4, 5.0, 9.5], dtype=np.float32)] 

condition = np.array([True, False, False, True])
data = np.array(data)[:,condition]

вывод:

[4.  4. ]
 [6.6 9.5]]

Answer 2

Если у вас есть список, выполните:

import numpy as np

data = [np.array([4,2,3,4], dtype=np.uint16),
        np.array([6.6, 7.4, 5.0, 9.5], dtype=np.float32)]

condition = np.array([True, False, False, True])


result = [e[condition] for e in data]
print(result)