使用np.where在2D数组中查找匹配的行

时间:2017-02-02 01:15:33

标签: python arrays python-2.7 numpy

我想知道如何将np.where用于2D数组

我有以下数组:

arr1 = np.array([[ 3.,  0.],
                 [ 3.,  1.],
                 [ 3.,  2.],
                 [ 3.,  3.],
                 [ 3.,  6.],
                 [ 3.,  5.]])

我想找到这个数组:

arr2 = np.array([3.,0.])

但是当我使用np.where()时:

np.where(arr1 == arr2)

它返回:

(array([0, 0, 1, 2, 3, 4, 5]), array([0, 1, 0, 0, 0, 0, 0]))

我无法理解它的含义。有人可以帮我解释一下吗?

1 个答案:

答案 0 :(得分:5)

您可能希望所有行等于arr2

>>> np.where(np.all(arr1 == arr2, axis=1))
(array([0], dtype=int64),)

这意味着第一行(第零个索引)匹配。

你的方法的问题是numpy广播数组(用np.broadcast_arrays可视化):

>>> arr1_tmp, arr2_tmp = np.broadcast_arrays(arr1, arr2)
>>> arr2_tmp
array([[ 3.,  0.],
       [ 3.,  0.],
       [ 3.,  0.],
       [ 3.,  0.],
       [ 3.,  0.],
       [ 3.,  0.]]) 

然后进行元素比较:

>>> arr1 == arr2
array([[ True,  True],
       [ True, False],
       [ True, False],
       [ True, False],
       [ True, False],
       [ True, False]], dtype=bool)
然后

np.where会为您提供每个True的坐标:

>>> np.where(arr1 == arr2)
(array([0, 0, 1, 2, 3, 4, 5], dtype=int64),
 array([0, 1, 0, 0, 0, 0, 0], dtype=int64))
#       ^---- first match (0, 0)
#          ^--- second match (0, 1)
#             ^--- third match (1, 0)
#  ...

这意味着(0, 0)(第一行左项)是第一个True,然后是0, 1(第一行右项),然后是1, 0(第二行,左项) ),....

如果沿第一个轴使用np.all,则会获得完全相等的所有行:

>>> np.all(arr1 == arr2, axis=1)
array([ True, False, False, False, False, False], dtype=bool)

如果保持尺寸,可以更好地可视化:

>>> np.all(arr1 == arr2, axis=1, keepdims=True)
array([[ True],
       [False],
       [False],
       [False],
       [False],
       [False]], dtype=bool)