如何在不使用循环的情况下用列/行索引替换 (0,1) numpy 数组的非零元素？

Question

I have a numpy array of 0's and 1's.

def random_array(p):
    return np.random.choice(2, 6, p=[p, 1-p])
my_matrix = np.array([random_array(j) for j in np.random.uniform(0.3, 1.0, 4)])

I can get indices of all num-zero elements as np.nonzero(my_matrix) . May you help me to allocate all these indices to the column number (even column indices will work, where we start from 0 rather 1) in which it is present. For example,

array([[0, 2, 0, 4, 5, 6],
       [0, 0, 0, 4, 5, 0],
       [1, 2, 0, 4, 0, 0],
       [0, 2, 3, 4, 5, 6]])

Here, all 1's have been replaced with the column number. Thus, this is for a matrix whose all non-zero elements were 1. If you can find the indices of column then that would also be fibulas because, I can get the same by adding 1.

Note: I do not wish to use any loop for this task.

Answer 1

You can multiply your array with another array of the same size containing the corresponding row/column index:

## Dummy data
#  Array size
s = (6,4)
#  Axis along which we need to calculate the index:
a = 0
#  Random binary array
x = np.random.rand(*s).round()

#  Get the index along one axis using broadcasting (starting with 1)
x = x*(np.expand_dims(range(s[a]),len(s)-a-1)+1)

Answer 2

If I understand well, you want to replace all no zero elements with their column indices (starting from 1 instead of 0) right? Then you can do it like:

idx = np.nonzero(my_matrix)
my_matrix[idx[0], idx[1]] = idx[1]+1

Answer 3

In [169]: def random_array(p):
     ...:     return np.random.choice(2, 6, p=[p, 1-p])
     ...: my_matrix = np.array([random_array(j) for j in np.random.uniform(0.3, 1.0, 4)])
In [170]: my_matrix
Out[170]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 1, 1, 0, 0, 0],
       [0, 0, 1, 1, 0, 1],
       [1, 1, 0, 0, 1, 0]])

Just multiply the range index. By broadcasting (6,) arange is fine for columns:

In [171]: np.arange(1,7)*my_matrix
Out[171]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 2, 3, 0, 0, 0],
       [0, 0, 3, 4, 0, 6],
       [1, 2, 0, 0, 5, 0]])

for rows

In [172]: np.arange(1,5)[:,None]*my_matrix
Out[172]: 
array([[0, 0, 0, 0, 0, 0],
       [0, 2, 2, 0, 0, 0],
       [0, 0, 3, 3, 0, 3],
       [4, 4, 0, 0, 4, 0]])