Question

我有一个numpy数组（2个元素列表的列表）a在下面给出，我有一个我希望寻找的2个元素[30.94, 0.]的列表。

当我做了以下操作时，我没有得到理想的结果。为什么呢？

import numpy as np
a = np.array([[  5.73,   0.  ],
              [ 57.73,  10.  ],
              [ 57.73,  20.  ],
              [ 30.94,   0.  ],
              [ 30.94,  10.  ],
              [ 30.94,  20.  ],
              [  4.14,   0.  ],
              [  4.14,  10.  ]])

np.where(a==np.array([30.94, 0.]))

但是我得到了

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

这不是真的。

Answer 1

正如迪瓦卡所暗示的那样，a == np.array([30.94, 0.])并不是你所期望的。数组是广播的，比较是按元素进行的。结果如下：

array([[False,  True],
       [False, False],
       [False, False],
       [ True,  True],
       [ True, False],
       [ True, False],
       [False,  True],
       [False, False]], dtype=bool)

但是，我们可以通过np.all获得我们想要的内容：

>>> np.all(a==np.array([30.94, 0.]), axis=-1)
array([False, False, False,  True, False, False, False, False], dtype=bool)
>>> np.where(_)
(array([3]),)

所以你可以看到第3行符合预期。请注意，使用带有浮点数的==的常见注意事项将适用于此处。

Answer 2

又一个解决方案，但请注意，这会比Dietrich's solution慢一点，特别是对于大型阵列。

In [1]: cond = np.array([30.94, 0.])
In [2]: arr = np.array([[  5.73,   0.  ],
                       [ 57.73,  10.  ],
                       [ 57.73,  20.  ],
                       [ 30.94,   0.  ],
                       [ 30.94,  10.  ],
                       [ 30.94,  20.  ],
                       [  4.14,   0.  ],
                       [  4.14,  10.  ]])

In [3]: [idx for idx, el in enumerate(arr) if np.array_equal(el, cond)]
Out[3]: [3]

numpy.where函数有什么问题？

2 个答案: