Numpy 2d 数组，获取指定列索引等于 1 的行的索引

Question

我有一个像这样的二维 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。同样，@987654324 @ 应该返回 [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]

Answer 2

解决方案

试试这个。

代码 - 逻辑

• 查找给定列 (target_col_index) 具有 1 的行和列索引。
• 仅选择列索引与target_col_index 匹配的行。
• 通过检查哪些行索引的结果与target_col_index 中的列数匹配，从而缩小结果范围。

import numpy as np

target_col_index = [1,2]
target_row_index = get_row_index(a, target_col_index)
print(target_row_index)

## Output
# [1,3]

## Other cases tested
# test_col_indexes = [ [0,1], [0,2], [0,3], [1,2], [1,3], [2,3], [0,1,3], [1,2,3] ]
# returned_row_indexes = [ [], [0], [], [1,3], [2,3], [3], [], [3] ]

代码 - 自定义函数

def get_row_index(arr, target_col_index=None):
    if target_col_index is None:
        return None
    else:
        row_index, col_index = np.where(arr==1)
        result = row_index[np.isin(col_index, target_col_index)]
        rows, counts = np.unique(result, return_counts=True)
        target_row_index = rows[counts==len(target_col_index)]
        return target_row_index

虚拟数据

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