Fast Sorting in Numpy

• Using np.sort and np.argsort

# In[1]
x=np.array([2,1,4,3,5])
np.sort(x)

# Out[1]
array([1, 2, 3, 4, 5])

• You can also use sort method.

# In[2]
x.sort()
print(x)

# Out[2]
[1 2 3 4 5]

• A related function is argsort, which instead returns the indices of the sorted elements.

# In[3]
x=np.array([2,1,4,3,5])
i=np.argsort(x)
print(i)
print(x[i])

# Out[3]
[1 0 3 2 4]
[1,2,3,4,5]

Sorting Along Rows and Columns

# In[4]
rng=np.random.default_rng(seed=42)
X=rng.integers(0,10,(4,6))
print(X)

# Out[4]
[[0 7 6 4 4 8]
 [0 6 2 0 5 9]
 [7 7 7 7 5 1]
 [8 4 5 3 1 9]]

---

# In[5]
np.sort(X,axis=0)

# Out[5]
array([[0, 4, 2, 0, 1, 1],
       [0, 6, 5, 3, 4, 8],
       [7, 7, 6, 4, 5, 9],
       [8, 7, 7, 7, 5, 9]])

---

# In[6]
np.sort(X,axis=1)

# Out[6]
array([[0, 4, 4, 6, 7, 8],
       [0, 0, 2, 5, 6, 9],
       [1, 5, 7, 7, 7, 7],
       [1, 3, 4, 5, 8, 9]])

Partial Sorts : Partitioning

• np.partition : takes an array and a number K; the result is a new array with the smallest K values to the let of the partition, and the remaining values to the right, in arbitrary order.

# In[7]
x=np.array([7,2,3,1,6,5,4])
np.partition(x,3)

# Out[7]
array([2, 1, 3, 4, 6, 5, 7])

• The first three values in the resulting array are the three smallest in the array, and the remaining array positions contain the remaining values.

# In[8]
np.partition(X,2,axis=1)

# Out[8]
array([[0, 4, 4, 7, 6, 8],
       [0, 0, 2, 6, 5, 9],
       [1, 5, 7, 7, 7, 7],
       [1, 3, 4, 5, 8, 9]])