在Python中打印包含非零元素的位置

Question

下面的代码打印行号solution1，其中至少有一个非零元素。然而，与这些行号相对应，我如何也打印出哪些位置有非零元素solution2，如预期输出所示。例如，行1在位置[1,3,4,6]处有非零元素，行2在位置[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)

当前输出为

solution 1 = [1,2]

期望输出为

solution 1 = [1,2]
solution 2 = [[1,3,4,6],[0,2,3,5]]

Answer 1

首先使用np.where查找所有非零值的坐标，然后按行分割y索引:

idx, idy = np.where(A)
np.split(idy, np.flatnonzero(np.diff(idx) != 0) + 1)
# [array([1, 3, 4, 6], dtype=int32), array([0, 2, 3, 5], dtype=int32)]