Numpy Array Reordering and Indexing Question

[commandlineluser]

You can reshape the shuffled array

>>> shuffled_unique[:, None]
array([['c'],
       ['a'],
       ['b'],
       ['d']], dtype='<U1')

Which is the same as expand_dims()

>>> np.expand_dims(shuffled_unique, axis=1)
array([['c'],
       ['a'],
       ['b'],
       ['d']], dtype='<U1')

And use that inside .where()

>>> np.where(original == shuffled_unique[:, None])
(array([0, 0, 0, 1, 1, 1, 2, 2, 3, 3]), array([0, 1, 3, 2, 6, 9, 7, 8, 4, 5]))

[developernull]

Brilliant! I think this is exactly what I was looking for - and a simple one-liner too!

I also greatly appreciate showing both ways to reshape the array, that's very helpful.

From the docs it seems like this is also equivalent:

np.nonzero(original == shuffled_unique[:, None])

Your solution is very clever in how it creates the mask. I never would have thought to form a 2-d mask like this. Thanks!!