python - Numpy:快速设置数组元素

Question

我刚刚遇到了我的一部分代码的问题，我不确定我做错了什么......

我只是有一个大数据立方体，想将 z 轴上的最大值更改为其他数字:

import numpy as np
from time import time

x, y, z = 100, 100, 10
a = np.arange(x*y*z).reshape((z, y, x))

t = time()
a[np.argmax(a, axis=0)] = 1
print(time() - t)

这大约需要 0.02 秒，对于这么小的数组来说有点慢，但还可以。我的问题是，我需要对大至 (32, 4096, 4096) 的数组执行此操作，但我没有耐心用上面的代码完成此操作……效率太低了，但实际上应该很快!我在设置数组元素时做错了吗？

Answer 1

你基本上是用一个包含数字的 numpy 数组来索引你的 numpy 数组。我认为这就是它如此缓慢的原因(而且我不确定它是否真的按照您的意愿进行)。

如果您创建一个 bool 型 numpy 数组并将其用作切片，它会快几个数量级。

例如:

pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
pos_max_indices = np.arange(a.shape[0]).reshape(10,1,1) == pos_max
a[pos_max_indices] = 1

比原版快 20 倍，性能相同。