Python中识别数组每行非自身位置的1并生成列表

我需要处理一个NumPy数组A，为每行生成一个列表，其中包含该行里除行号对应位置外所有值为1的元素索引。例如行A[0]中，要提取索引2,3,5，排除索引0。

当前代码 #

import numpy as np

A=np.array([[1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0],
       [1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0],
       [1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1],
       [0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
       [0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1],
       [0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1]])

output = []
for i, row in enumerate(A):
    ones_indices = np.where(row == 1)[0]
    other_rows = np.arange(A.shape[0])
    other_rows = np.delete(other_rows, i)
    output.append([[i], other_rows.tolist()])

print(output)

当前输出 #

[[[0], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]], [[1], [0, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]], [[2], [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11]], [[3], [0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11]], [[4], [0, 1, 2, 3, 5, 6, 7, 8, 9, 10, 11]], [[5], [0, 1, 2, 3, 4, 6, 7, 8, 9, 10, 11]], [[6], [0, 1, 2, 3, 4, 5, 7, 8, 9, 10, 11]], [[7], [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11]], [[8], [0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 11]], [[9], [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 11]], [[10], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 11]], [[11], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]]

期望输出 #

[[[0], [2,3,5]], [[1], [3,4,6]], [[2], [0,3,5]], [[3], [0, 1, 2, 4, 5, 6]], [[4], [1,3,6]], [[5], [0,2,3,7,8,10]], [[6], [1,3,4,8,9,11]], [[7], [5,8,10]], [[8], [5,6,7,9,10,11]], [[9], [6,8,11]], [[10], [5,7,8]], [[11], [6,8,9]]]

解决方案 #

当前代码的问题在于，错误生成了排除当前行号的所有行号列表，而非提取该行中值为1且不等于行号的索引。修正后的代码如下：

import numpy as np

A=np.array([[1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0],
       [1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0],
       [1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, 1],
       [0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1],
       [0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1],
       [0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 0],
       [0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1]])

output = []
for i, row in enumerate(A):
    # 获取当前行所有值为1的索引
    ones_indices = np.where(row == 1)[0]
    # 过滤掉等于当前行号i的索引，转为列表
    filtered_indices = ones_indices[ones_indices != i].tolist()
    # 按期望格式添加到输出
    output.append([[i], filtered_indices])

print(output)

代码说明 #

1. 遍历数组每一行，同时记录行号i；
2. 用np.where(row == 1)[0]获取当前行所有值为1的索引；
3. 通过布尔索引ones_indices != i筛选掉等于行号的索引，再转为列表；
4. 按照[[行号], [筛选后的索引]]的格式添加到输出列表中。

执行上述代码后，输出将与期望结果完全一致。

内容的提问来源于stack exchange，提问作者KeplerNick123