Numpy - 在二維數組中查找一維數組中元素的出現

Question

我正在努力修復以下代碼：

import numpy as np

arr1 = np.array([2,15,132,135,136])
arr2 = np.array([[0,0,0,0,0], [2,15,132,307,406], [15,132,135,136,148]])

s1 = np.sum(np.all(arr1 == arr2[:, None], axis=1), axis=1)
print(s1)

我期待的結果是[0 3 4] ，以表明：

1. 在 arr2 中找不到行索引 0 的匹配項
2. 在 arr2 中為行索引 1 找到了三個匹配項
3. 在 arr2 中為行索引 2 找到了四個匹配項

我得到的實際結果是[0 3 0]

有人可以幫助我理解這一點並幫助我修改代碼以產生所需的結果嗎？

Answer 1

>>> (arr1 == arr2[...,None]).sum(axis=(1, 2))
array([0, 3, 4])

您想為arr2的最后一個軸進行廣播。 通過使用冒號而不是省略號，您可以在第一個軸之后添加新軸。

>>> arr2[:,None].shape
(3, 1, 5)
>>> arr2[...,None].shape   # arr2[:,:,None] would also work
(3, 5, 1)