numpy：获取每行满足条件的索引

Question

我有一个如下的数组：

In [70]: x
Out[70]: 
array([[0, 1, 2],
       [3, 4, 5]])

我正在尝试获取每行满足条件的索引，例如x > 1。

期望的输出类似于([2], [0, 1, 2])

我尝试过numpy.where，numpy.nonzero，但它们给出了奇怪的结果。

- cosmosa

Answer 1

一种方法是 -

r,c = np.where(x>1)
out = np.split(c, np.flatnonzero(r[1:] > r[:-1])+1)

样例运行 -

In [140]: x
Out[140]: 
array([[0, 2, 0, 1, 1],
       [2, 2, 1, 2, 0],
       [0, 2, 1, 1, 0],
       [1, 0, 0, 2, 2]])

In [141]: r,c = np.where(x>1)

In [142]: np.split(c, np.flatnonzero(r[1:] > r[:-1])+1)
Out[142]: [array([1]), array([0, 1, 3]), array([1]), array([3, 4])]

或者，我们可以在最后一步使用np.unique，如下所示-

np.split(c, np.unique(r, return_index=1)[1][1:])

- Divakar

---

把它封装在一个函数中，取个花哨的名字，假装它是内置的。就这样，它变得“更简单”了 :) - Divakar