使用 numpy 优化数组计算

Question

我尝试在 Python 中优化以下算法：

我事先定义了三个逻辑矩阵（X、Y 和 Z）。每个都有 N 列。我还定义了两个非常大的数字 N 和 A，以及一个名为 tot_list 的任意数字数组，其长度为 A。我想遍历 X、Y 和 Z，找到每个元素的元素为 1 的位置，然后更改它列表 A_list 中的对应值到随机浮点数。完成此操作后，我将 tot_list 更新为 A_list。下面是实际代码：

for i in range(0,N):
    A_list = np.ones([1,A])
    
    A_list[0][np.where(X[:,i]==1)[0]] = np.random.random()
    A_list[0][np.where(Y[:,i]==1)[0]] = np.random.random()
    A_list[0][np.where(Z[:,i]==1)[0]] = np.random.random()
    
    tot_list = tot_list+A_list

上面给出的代码完全符合我的要求，但运行速度非常慢。我必须处理巨大的数字（即，A 大约为 10^6，N 大约为 10^2），并且我必须尽可能优化上述内容。我已经尝试过摆脱 for 循环，但我不确定如何实现这一点。欢迎提出任何建议。

Answer 1

好的，如果没有示例输入和输出，很难知道问题所在，但让我尝试回答，如果不正确，我可以编辑。

正如我所见，您正在循环遍历 X、Y 和 Z 的每一行，并找到每个为 1 的元素。您可以一次对单个数组的每个元素执行此操作，只需执行例如X == 1.

然后您找到 X、Y、Z 等于 1 的位置并将其设置为随机数，否则将其设置为 1（根据您对 A_list 的定义）。所以我们可以找到所有等于 1 的数组元素中的OR，例如

equal_one = (X == 1) | (Y == 1) | (Z == 1)

那么我们可以使用np.where的三参数版本将等于1的值设置为随机数，将不等于1的值设置为1。

out = np.where(equal_one, np.random.random(), 1)

最后，我们沿一个轴对数组的每一行求和。我认为从你的例子这相当于

tot_list = out.sum(axis=1)

把它们放在一起

tot_list = np.where(
    (X == 1) | (Y == 1) | (Z == 1),
    np.random.random(),
    1
).sum(axis=1)