Printing locations containing non-zero elements in Python

Question

The following code prints row numbers solution1 which have at least one non-zero element. However, corresponding to these row numbers, how do I also print which locations have non-zero elements solution2 as shown in the expected output.? For instance, row 1 has non-zero elements at locations [1,3,4,6], row 2 has non-zero elements at locations [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)

The current output is

solution 1 = [1,2]

The expected output is

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

Asked By: AEinstein

||

Source

Answer 1

Use np.where to find all coordinates for non zero values first, and then split y index by rows:

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)]

Answered By: Psidom

Printing locations containing non-zero elements in Python

Question:

Answers: