python - Numpy 二维数组，获取指定列索引等于 1 的行的索引

Question

我有一个像这样的 2d numpy 数组，它只会有 0、1 个值。

a = np.array([[1, 0, 1, 0],  # Indexes 0 and 2 == 1
             [0, 1, 1, 0],   # Indexes 1 and 2 == 1
             [0, 1, 0, 1],   # Indexes 1 and 3 == 1
             [0, 1, 1, 1]])  # Indexes 1, 2, and 3 == 1

我想做的是获取一对列索引传递的每一行的索引都等于1。

例如，如果执行此操作的函数是 get_rows，则 get_rows(a, [1, 3]) 应该返回 [2, 3]，因为行在索引 2 和 3 处，列索引 1 和 3 等于 1。类似地，get_rows(a, [1, 2]) 应该返回 [1, 3]。

我知道如何在 Pandas 数据框中执行此操作，但我想坚持为此使用纯 numpy。我尝试以某种形式使用 np.where

np.where( ((a[i1 - 1] == 1) & (a[i2 - 1] == 1) ))

但这似乎并没有给我我想要的东西，并且不适用于不同数量的传递索引。

Answer 1

我想你正在寻找这个:

col_idx = [1, 2]
np.where(a[:,col_idx].all(axis=1))[0]

您可以使用任何要传递给它的列索引。使用 np.where 提取列并搜索其中全为 1 的行是非常不言自明的。

编辑:根据@Mad Physicist 的建议，这是另一个类似的解决方案:

np.flatnonzero(a[:,col_idx].all(axis=1))

输入的输出示例:

[1 3]