Changing the format of the output in Python

Question

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

Asked By: modishah123

||

Source

Answer 1

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

Answered By: Mattravel

Answer 2

If you always have the same number of 1s in the array, you can just combine where and a reshape:

C = np.where(A==1)[1].reshape(A.shape[0], -1)

NB. if you don’t have the same number of 1s, then you can’t build a non-object array.

Output:

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

Answered By: mozway

Changing the format of the output in Python

Question:

Answers: