Python调整输出格式：提取numpy数组行中1的位置并转换结构

方法1：改进循环实现 #

遍历数组每行，提取1的索引并整理成目标格式：

import numpy as np

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

index_list = []
for row in A:
    # 提取每行中1的索引数组（np.where返回元组，取第一个元素）
    idx = np.where(row == 1)[0]
    index_list.append(idx)

# 将索引数组堆叠为二维数组，再包装成列表
result = [np.vstack(index_list)]
print(result)

输出结果：

[array([[1, 2, 3],
       [0, 2, 3],
       [0, 1, 3],
       [0, 1, 2]], dtype=int64)]

方法2：向量化操作（更高效） #

利用Numpy的向量化特性直接处理整个数组，避免循环：

import numpy as np

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

# 获取所有值为1的位置的行列索引
rows, cols = np.where(A == 1)
# 按行数reshape为二维数组，再包装成列表
result = [cols.reshape(A.shape[0], -1)]
print(result)

输出结果与方法1完全一致。

原代码问题解析 #

原代码中np.where(A[i]==1)返回的是包含一维数组的元组（Numpy的where函数统一返回索引元组，即使处理一维输入），直接append到列表后得到的是元组的列表。我们需要提取元组内的索引数组，再将这些数组堆叠成二维数组，最终包装为符合要求的列表格式。

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