在Python中获取numpy数组每行值为1的列索引并完成格式转换

我来帮你一步步实现需求，分两个阶段完成：先提取每行中值为1的列索引，再进行格式转换。

首先先定义你的numpy数组：

import numpy as np

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

第一步：生成output1（获取每行1的列索引） #

我们可以通过遍历数组的每一行，结合np.where()函数定位值为1的元素位置，再转成列表格式：

output1 = [np.where(row == 1)[0].tolist() for row in p]
print(output1)  # 输出结果: [[1, 3], [], [0], [2]]

这里np.where(row == 1)[0]会返回该行中所有值为1的元素的列索引（以numpy数组形式），.tolist()则把它转换成普通Python列表，空行自然会生成空列表。

第二步：生成output2（索引+1，空行替换为[0]） #

基于output1，我们可以用循环或者列表推导式来处理每个子列表：

方式一：使用普通循环（更直观）

output2 = []
for sublist in output1:
    if not sublist:
        # 如果子列表为空，添加[0]
        output2.append([0])
    else:
        # 否则每个索引值加1
        output2.append([x + 1 for x in sublist])
print(output2)  # 输出结果: [[2, 4], [0], [1], [3]]

方式二：使用列表推导式（更简洁）

如果喜欢一行代码搞定，也可以写成：

output2 = [[0] if not sub else [x + 1 for x in sub] for sub in output1]

这样同样能得到你想要的output2结果。

备注：内容来源于stack exchange，提问作者Starwinwin