I have an array A. I am identifying locations on each row of A which has element 1. However, I want to change the format of the output. I present the current and expected outputs.

import numpy as np

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

C=[]
for i in range(0,len(A)): 
    B=np.where(A[i]==1)
    C.append(B)
print(C)

The current output is

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

The expected output is

[(array([[1, 2, 3],[0, 2, 3],[0, 1, 3],[0, 1, 2]])]

>Solution :

You can do the following:

import numpy as np

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

C=[]
for i in range(0,len(A)): 
    B=np.where(A[i]==1)
    C.append(B[0].tolist())
C = [np.array(C)]
print(C)

Output:

[array([[1, 2, 3],
       [0, 2, 3],
       [0, 1, 3],
       [0, 1, 2]])]