如何有效地将条件应用于numpy数组的索引？

debugcn 发表于 Dev

拉胡尔·维什瓦卡玛（Rahul Vishwakarma）

我有一个2D NumPy数组，我想为该数组设置值，前提是其索引满足特定条件。

我可以通过以下方式使用for循环：

import numpy as np

new_a = np.ones((5,10), dtype=np.float32)
for i in range(new_a.shape[0]):
    for j in range(new_a.shape[1]):
        if (np.nan_to_num(i/np.nan_to_num(j))) >= new_a.shape[0]/new_a.shape[1]:               #(This is the condition, which I may change a little as needed)
            new_a[i, j] = 0
            
print(new_a)

''' Output:                                                                                    (This gives a upper triangular matrix)
[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]]
'''

但是问题是我想对尺寸约为的大型2D数组执行此操作(10000, 20000)。因此for循环将非常慢。我如何使用NumPy或任何其他库有效地（在更少的时间内）做到这一点？

注意：我不希望解决方案创建对角矩阵（因为我想将代码应用于许多不同的条件）。我正在寻找“有效地将条件应用于numpy数组索引的解决方案”（比使用for循环要快得多的方法）的解决方案。

内森

您可以像这样获得数组索引：

import numpy as np

new_a = np.ones((5,10), dtype=np.float32)
indices = np.indices(new_a.shape)
y_indices = indices[0]
x_indices = indices[1]

要获取您的特定比较成立的索引，您可以：

locations = np.nan_to_num(indices[0] / indices[1]) >= new_a.shape[0] / new_a.shape[1]

要应用此功能，只需：

new_a[locations] = 0
print(new_a)

归来

[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
 [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]]

本文收集自互联网，转载请注明来源。

如有侵权，请联系[email protected] 删除。