Python/Numpy：如何避免存储布尔索引子数组？

Question

在使用Numpy时，我经常需要使用布尔索引来访问数组的各个部分。为了便于阅读和键入，我经常将这些子数组存储到新变量中，例如：

n = 10000
X = np.random.rand((n, n))
W = np.random.random_integers(0, 1, n)
X0 = X[W==0]
X1 = X[W==1]

然而，当我处理越来越大的数据集时，这看起来非常浪费。在这种情况下，推荐的做法是什么？我应该只写（在上面的例子中）X[W==0]和X[W==1]吗？你知道吗

Answer 1

如果要索引的元素数量有限，则可能需要将索引存储为索引列表。为此，numpy.nonzero非常有用。但是，如果要索引的元素数很大，则可以通过存储布尔数组（每个元素1字节）来使用较少的内存。你知道吗

因此，有四种可能性：

1. 存储布尔数组
2. 存储非零索引
3. 总是分开比较
4. 屏蔽阵列

从内存存储的角度来看，备选方案1为每个维度的每个索引元素占用8字节。（当然，可以通过使用平面索引来避免“每个维度”）布尔方法每个元素需要1个字节，因此如果布尔表中有超过1/8的元素带有True，那么它是更节省空间的解决方案。解#3可能与布尔解占用相同的空间。你知道吗

（我对NumPy的内部结构了解不多，所以不能对蒙面阵列说太多。我怀疑它们的行为类似于布尔索引。）

从表现上看，情况也差不多。如果有很多元素要挑选，布尔解决方案是有效的，但是如果只有很少的元素，那么索引解决方案更好。你知道吗

只是想给出一些基准测试的想法：

import numpy as np
import time

def create_indices(prob):
    data = np.random.random(100000000) < prob
    return data, np.nonzero(data)

def bool_index(data):
    return data[data]

def list_index(data, indices):
    return data[indices]

通过使用不同概率的timeit，结果是：

    p    boolean   list
   0.01   0.206    0.012
   0.10   0.415    0.099
   0.20   0.405    0.146
   0.50   0.786    0.373
   0.75   0.539    0.555
   1.00   0.214    0.723

这实际上非常有趣：当一半的元素是True时，使用布尔索引是最糟糕的。使用列表索引的行为与预期一致。你知道吗

不能把这一基准视为全部真理。可能是要索引的数组的类型改变了情况（这里的bool与uint8相同），等等。但是，在大多数情况下，性能方面的列表索引似乎非常好。你知道吗