在numpy中保持维度的数组索引列表。

Question

我有一个arryas的列表。

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数组，并像@Quang在评论中提到的那样利用numpy索引。

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

如果你有一个list，就做。

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)