我所拥有的:
import numpy as np
np.random.seed(42)
dlen = 250000
data = np.random.rand(dlen, 3, 3)
mask = np.random.choice([0, 1, 2], dlen)
我想要的是:
[[0.37454012 0.95071431 0.73199394],
[0.83244264 0.21233911 0.18182497],
[0.13949386 0.29214465 0.36636184],
[0.94888554 0.96563203 0.80839735],
[0.44015249 0.12203823 0.49517691],
....
(250000, 3)
我尝试使用的是:
data[:,mask,:]
{MemoryError}Unable to allocate 1.36 TiB for an array with shape (250000, 250000, 3) and data type float64
是什么给出了正确的结果,但看起来很奇怪:
data[np.arange(data.shape[0]), mask, :]
那么正确的使用方法是什么呢?
upd:掩码应选择具有指定索引的列。形状为[2,3,3]的数组示例:
array = [[[5 6 7], [7 8 9], [2 3 4]],
[[2 1 0], [7 6 5], [7 6 5]]]
mask = [1 0]
result = [[7 8 9],
[2 1 0]]
1条答案
按热度按时间6qftjkof1#
这是因为它是一个多维索引数组
当我在这里提到术语掩码时,我想到了布尔索引。您的整数掩码可以转换为布尔掩码,以您想要的方式使用它。
您可以使用第二个维度长度使其更为通用:
如果控制掩码的构造,则可能希望将其构造为布尔数组。
如果这是新代码,您可能希望升级到兼容版本的numpy并使用newrandom生成器。