如何在沒有for循環的情況下在測試數組的每個元素上使用numpy？

Question

我想在例程中使用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對函數進行矢量化可能會更容易，但是我感覺這對於問題來說過於np.vectorize過正，因此必須有一個更簡單的解決方案。 任何指導將不勝感激？

Answer 1

這是使用.outer和np.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,:]

因此，這是一個二test_elem == data數組，其行等於for循環中出現的布爾數組test_elem == data 。

where返回來自此兩個食指的陣列，每一個坐標。 x是瞬息萬變的協調和等於值1D where你for循環的回報，但都粘在一起成一個長的矢量。 y是慢速調整坐標，它的值是有序的，可以用於對x進行分組/拆分。 searchsorted可能不是最有效的方法，但它很簡單並且可以正確處理空行。