如何将二进制numpy.ndarray转换为嵌套列表（Python）

需求说明 #

将二进制numpy矩阵转换为嵌套列表，每个子列表对应矩阵一行中值为1的元素的列索引。

示例输入输出 #

输入矩阵（numpy.ndarray）：

import numpy as np

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

目标输出列表：

[[0,2,4],[2,3,4],[0,2,3],[2,4],[0,1,4],[0,2,3,4],[2,4]]

实现方法 #

方法1：结合numpy.where与列表推导式

利用numpy的where函数快速定位值为1的元素索引，再转换为列表：

result = [np.where(row == 1)[0].tolist() for row in matrix]

方法2：纯列表推导式（无需额外numpy函数）

遍历每行的元素与索引，筛选出值为1的索引：

result = [[idx for idx, val in enumerate(row) if val == 1] for row in matrix]

两种方法都能得到目标嵌套列表，可根据习惯选择。

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