Python中获取含非零元素的数组行号及对应非零元素位置

现有Python代码可打印出NumPy数组A中至少包含一个非零元素的行号solution1，当前输出为solution 1 = [1,2]。现需同时输出这些行中各自非零元素的列位置solution2，期望输出为solution 1 = [1,2]，solution 2 = [[1,3,4,6],[0,2,3,5]]。

原代码如下：

import numpy as np

A=np.array([[  0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 212.13245959,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 216.08166277, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]])


solution1 = []
for idx, e in enumerate(A):
    if any(e): 
        solution1.append(idx)
print("solution 1 =",solution1)

修改后的代码 #

import numpy as np

A=np.array([[  0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [  0.        , 423.81345923,   0.        , 407.01354328,
        419.14952534,   0.        , 212.13245959,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ],
       [402.93473651,   0.        , 216.08166277, 407.01354328,
          0.        , 414.17017965,   0.        ,   0.        ,
          0.        ,   0.        ,   0.        ,   0.        ]])


solution1 = []
solution2 = []
for idx, row in enumerate(A):
    # 获取当前行非零元素的列索引并转为列表
    non_zero_cols = np.where(row != 0)[0].tolist()
    if non_zero_cols:
        solution1.append(idx)
        solution2.append(non_zero_cols)
print("solution 1 =", solution1)
print("solution 2 =", solution2)

改动说明 #

• 新增solution2列表，用于存储每行非零元素对应的列索引
• 使用np.where(row != 0)[0]定位当前行的非零元素列位置，通过.tolist()转换为Python列表
• 当该行存在非零元素时，同步将行号加入solution1，列索引列表加入solution2

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