如何在没有for循环的情况下在测试数组的每个元素上使用numpy?

时间:2019-06-15 10:26:13

标签: python-3.x numpy for-loop vectorization where

我想在例程中使用numpy函数而不使用for循环。考虑下面的示例:

import numpy as np

data = np.linspace(1, 10, 10).astype(int) 
test_elements = np.array([1, 2])
for test_elem in test_elements:
    print(np.where(test_elem == data))

...

(array([0]),)
(array([1]),)

我还阅读了其他文章和numpy文档。常见的建议似乎是使用np.roll滚动test_elements,或使用其他技巧,例如大步迈进(我不完全理解)。我认为使用np.vectorize对函数进行向量化可能会更容易,但是我感觉这对于问题来说过于矫kill过正,因此必须有一个更简单的解决方案。任何指导将不胜感激?

1 个答案:

答案 0 :(得分:1)

这是使用.outernp.split的一种方法,我使示例变得更加有趣。

data = np.linspace(1, 5, 10).astype(int) 
test_elements = np.array([1, 2, 4, 6])
y, x = np.where(np.equal.outer(test_elements,data))
np.split(x, y.searchsorted(np.arange(1,test_elements.size)))
# [array([0, 1, 2]), array([3, 4]), array([7, 8]), array([], dtype=int64)]

一些解释:

np.equal.outer(test_elements,data)

相同
 test.elements[:,None] == data[None,:]

这是一个二维布尔数组,其行等于您的for循环中出现的布尔数组test_elem == data

where从中返回两个索引数组,每个坐标数组一个。 x是快速变化的坐标,它等于for循环的1d where返回的值,但所有这些都粘在一起形成一个长向量。 y是慢速引导坐标,它的值是有序的,可用于对x进行分组/拆分。 searchsorted可能不是最有效的方法,但它很简单并且可以正确处理空行。