numpy.where返回元组的目的是什么？

Question

当我运行此代码时：

import numpy as np
a = np.array([1, 2, 3, 4, 5, 6])
print(np.where(a > 2))

自然会得到一个索引数组，其中a > 2，即[2, 3, 4, 5]，但是我们得到：

(array([2, 3, 4, 5], dtype=int64),)

即具有空第二个成员的元组。

然后，要获得的“自然”答案numpy.where，我们必须这样做：

np.where(a > 2)[0]

这个元组有什么意义？在什么情况下有用？

注意：我在这里仅谈论用例numpy.where(cond)，但并不numpy.where(cond, x, y)存在（请参阅文档）。

Answer 1

numpy.where 返回一个元组，因为该元组的每个元素都引用一个维度。

考虑以下两个方面的示例：

a = np.array([[1, 2, 3, 4, 5, 6],
              [-2, 1, 2, 3, 4, 5]])

print(np.where(a > 2))

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

如您所见，元组的第一个元素指的是相关元素的第一个维度；第二个元素是指第二个维度。

这是一个numpy经常使用的约定。当您要求数组的形状时，您也会看到它，即一维数组的形状将返回一个包含1个元素的元组：

a = np.array([[1, 2, 3, 4, 5, 6],
              [-2, 1, 2, 3, 4, 5]])

print(a.shape, a.ndim)  # (2, 6) 2

b = np.array([1, 2, 3, 4, 5, 6])

print(b.shape, b.ndim)  # (6,) 1

Answer 2

从文档 np.where

如果只给出条件，则返回元组 condition.nonzero()，条件为 True 的索引

所以我们查看'np.nonzero'的文档

返回一个数组元组，a 的每个维度一个，包含该维度中非零元素的索引。a 中的值始终以行优先的 C 样式顺序进行测试和返回。相应的非零值可以通过以下方式获得：

因此，如何能这是非常有用的np.where/np.nonzero返回数组的一个元组？我认为这与索引多维数组有关。

从文档的示例中，如果我们有

y = np.arange(35).reshape(5,7)

我们可以做的

y[np.array([0,2,4]), np.array([0,1,2])]

选择y[0, 0], y[2, 1], y[4, 2]。

在这种情况下，如果索引数组具有匹配的形状，并且被索引的数组的每个维度都有一个索引数组，则结果数组与索引数组具有相同的形状，并且值对应于为每个数组设置的索引在索引数组中的位置。在此示例中，两个索引数组的第一个索引值都是 0，因此结果数组的第一个值是 y[0,0]。下一个值是 y[2,1]，最后一个值是 y[4,2]。

希望索引多维数组可以证明np.nonzero/np.where返回一个数组元组是正确的，以便以后可以使用它来选择元素。

Answer 3

为了一致性：元组的长度与输入数组的维数相匹配。

>>> np.where(np.ones((1)) > 0)
(array([0]),)
>>> np.where(np.ones((1,1)) > 0)
(array([0]), array([0]))
>>> np.where(np.ones((1,1,1)) > 0)
(array([0]), array([0]), array([0]))

使一维情况返回数组而不是元组会导致不均匀的返回类型。如果调用者代码正在处理任意形状的输入数据，那么程序员必须对返回值中的一维输入进行特殊情况处理。