argsort for a multidimensional ndarray

by

Solution:

>>> a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]
array([[1, 2, 3],
       [2, 8, 9]])

You got it right, though I wouldn’t describe it as cheating the indexing.

Maybe this will help make it clearer:

In [544]: i=np.argsort(a,axis=1)

In [545]: i
Out[545]: 
array([[1, 2, 0],
       [2, 0, 1]])

i is the order that we want, for each row. That is:

In [546]: a[0, i[0,:]]
Out[546]: array([1, 2, 3])

In [547]: a[1, i[1,:]]
Out[547]: array([2, 8, 9])

To do both indexing steps at once, we have to use a ‘column’ index for the 1st dimension.

In [548]: a[[[0],[1]],i]
Out[548]: 
array([[1, 2, 3],
       [2, 8, 9]])

Another array that could be paired with i is:

In [560]: j=np.array([[0,0,0],[1,1,1]])

In [561]: j
Out[561]: 
array([[0, 0, 0],
       [1, 1, 1]])

In [562]: a[j,i]
Out[562]: 
array([[1, 2, 3],
       [2, 8, 9]])

If i identifies the column for each element, then j specifies the row for each element. The [[0],[1]] column array works just as well because it can be broadcasted against i.

I think of 

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

as ‘short hand’ for j. Together they define the source row and column of each element of the new array. They work together, not sequentially.

The full mapping from a to the new array is:

[a[0,1]  a[0,2]  a[0,0]
 a[1,2]  a[1,0]  a[1,1]]

def foo(a):
    i = np.argsort(a, axis=1)
    return (np.arange(a.shape[0])[:,None], i)

In [61]: foo(a)
Out[61]: 
(array([[0],
        [1]]), array([[1, 2, 0],
        [2, 0, 1]], dtype=int32))
In [62]: a[foo(a)]
Out[62]: 
array([[1, 2, 3],
       [2, 8, 9]])

More Related Contents:

  1. Sorting arrays in NumPy by column
  2. Find elements of array one nearest to elements of array two
  3. Sort invariant for numpy.argsort with multiple dimensions
  4. Efficiently sorting a numpy array in descending order?
  5. Sorting a python array/recarray by column
  6. Performance of Pandas apply vs np.vectorize to create new column from existing columns
  7. How to access the ith column of a NumPy multidimensional array?
  8. Python/NumPy first occurrence of subarray
  9. Concatenating two one-dimensional NumPy arrays
  10. Numpy matrix to array
  11. Replacing Numpy elements if condition is met
  12. How to make a multidimension numpy array with a varying row size?
  13. How can I add new dimensions to a Numpy array?
  14. Force NumPy ndarray to take ownership of its memory in Cython
  15. How to normalize a NumPy array to within a certain range?
  16. subsampling every nth entry in a numpy array
  17. Immutable numpy array?
  18. Sort multidimensional array based on 2nd element of the subarray
  19. How to get a list of all indices of repeated elements in a numpy array
  20. Check if two 3D numpy arrays contain overlapping 2D arrays
  21. Numpy: Creating a complex array from 2 real ones?
  22. Repeat each values of an array different times
  23. How to create 2d array with numpy random.choice for every rows?
  24. What is the difference between np.array() and np.asarray()?
  25. Flattening a list of NumPy arrays?
  26. Fastest pairwise distance metric in python
  27. Why is Numpy much faster at creating a Zero array compared to replacing the values of an existing array with zeros?
  28. sort eigenvalues and associated eigenvectors after using numpy.linalg.eig in python
  29. Generate random array of floats between a range
  30. Numpy `ValueError: operands could not be broadcast together with shape …`

Leave a Comment