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相同）等。但是，在大多数情况下，性能方面的列表索引似乎非常好。