numpy：布尔索引和内存使用

Posted 2023-03-11

【中文标题】numpy：布尔索引和内存使用

【英文标题】：numpy: boolean indexing and memory usage

【发布时间】：2011-08-23 03:31:01

【问题描述】：

考虑以下numpy 代码：

A[start:end] = B[mask]

这里：

A 和 B 是列数相同的二维数组； start 和 end 是标量； mask 是一维布尔数组； (end - start) == sum(mask)。

上述操作原则上可以使用O(1)临时存储进行，将B的元素直接复制到A中。

这是实际发生的情况，还是numpy 为B[mask] 构造了一个临时数组？如果是后者，有没有办法通过重写语句来避免这种情况？

【参考方案1】：

线

A[start:end] = B[mask]

将 - 根据 Python 语言定义 - 首先计算右侧，生成一个包含选定行的新数组 B 并占用额外的内存。我知道避免这种情况的最有效的纯 Python 方法是使用显式循环：

from itertools import izip, compress
for i, b in izip(range(start, end), compress(B, mask)):
    A[i] = b

当然，这比原始代码的时间效率要低得多，但它只使用 O(1) 额外的内存。另请注意，itertools.compress() 在 Python 2.7 或 3.1 或更高版本中可用。

【讨论】：

当然，“产生一个包含 B 的选定行并占用额外内存的新数组”是不合理的吗？由B.__getitem__() 选择它想要返回的内容。例如，如果mask 是slice，则将返回代理（视图），并且不会进行复制。

@aix：根据 OP，mask 是一维布尔数组。我错过了什么吗？

@aix：哦，我明白了。语言定义的部分有点模棱两可。它只是指“首先评估右手边”的部分。

是的，我想我们彼此了解。

【参考方案2】：

使用布尔数组作为索引是一种花哨的索引，因此 numpy 需要进行复制。 如果您遇到内存问题，您可以编写一个 cython 扩展来处理它。

【讨论】：

+1 用于引入 Cython。它擅长的正是这种循环。