numpy.where 函数有什么问题？答案

【问题标题】：What is wrong with numpy.where function?numpy.where 函数有什么问题？
【发布时间】：2016-12-22 20:03:59
【问题描述】：

我有一个 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]))

这不是真的。

【问题讨论】：

问你一个问题 - 你对a==np.array([30.94, 0.]有什么期望？
行的索引应该是[3, 3]，列的索引应该是[0, 1]或类似的东西。
不，我的意思是没有np.where，只是：a==np.array([30.94, 0.]。我的想法是，如果你看看它的输出，你会有更好的画面，希望如此。
也许[0, 0, 0, 1, 0, 0, 0, 0] ?

标签： python numpy multidimensional-array numpy-ndarray

【解决方案1】：

正如 Divakar 所暗示的，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 行匹配，正如预期的那样。请注意，将== 与浮点数一起使用的常见注意事项将适用于此处。

【讨论】：

为了考虑浮动点数，我们可以使用np.isclose()，它允许broadcasting。
@kmario23 “broadcast”是“broadcast”的有效过去分词。
@DietrichEpp 我刚刚意识到这一点，但以ed 形式拥有它感觉更自然。感谢您指出，顺便说一句！

【解决方案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]

【讨论】：