Rearranging entries of rows in numpy array

Question

I am trying to transform transform a numpy array A into B without using a loop.

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

B=np.array([[1,2,3],
            [1,0,3],
            [0,2,0]])

So in each row, I want to reorder the entries using their value as the index. (i.e. in row 2, [1,3,0], 1 is the first entry, 3 is the third entry, and the 0 would fill in as the 2nd entry to make it [1,0,3].

I can do this for a single row so I could loop through the array, but I wanted to see if there was a way to do this without a loop. I know looping won't make a difference on small arrays like this, but I fear looping will create a bottleneck when doing this on large arrays (1m,1m).

Thank you!

Answer 1

Interesting question, +1.

In [28]:

import numpy as np
A=np.array([[1,2,3],
            [1,3,0],
            [2,0,0]])
In [29]:

B=np.zeros(A.shape, 'int64')+np.arange(1, A.shape[0]+1)
In [30]:

np.where(np.asarray(map(np.in1d, B, A)), B, 0)
Out[30]:
array([[1, 2, 3],
       [1, 0, 3],
       [0, 2, 0]])
In [31]:

%timeit np.where(np.asarray(map(np.in1d, B, A)), B, 0)
10000 loops, best of 3: 156 µs per loop

Answer 2

A different approach to get the same thing. Will probably be faster for large arrays:

>>> mask = A != 0
>>> rows, cols = A.shape
>>> idx = (A - 1 + (np.arange(rows)*cols)[:, None])[mask]
>>> B = np.zeros_like(A)
>>> B.ravel()[idx] = A[mask]
>>> B
array([[1, 2, 3],
       [1, 0, 3],
       [0, 2, 0]])

It converts the non-zero entries of A into indices in a flattened array, then uses those indices to copy the non-zero entries of A into their right positions in a flattened view of B.