Geting The K-smallest Values Of Each Column In Sorted Order Using Numpy.argpartition

Using np.argpartition, it does not sort the entire array. It only guarantees that the kth element is in sorted position and all smaller elements will be moved before it. Thus, the

Solution 1:

To use np.argpartition and maintain the sorted order, we need to use those range of elements as range(k) instead of feeding in just the scalar kth param -

idx = np.argpartition(myBigArray, range(num), axis=1)[:, :num]
out = myBigArray[np.arange(idx.shape[0])[:,None], idx]

Solution 2:

You can use the exact same trick that you used in the case of rows; combining with @Divakar's trick for sorting, this becomes

In [42]: num = 2

In [43]: myBigArray[np.argpartition(myBigArray, range(num), axis=0)[:num, :], np.arange(myBigArray.shape[1])[None, :]]
Out[43]: 
array([[ 1,  3,  2,  5,  7,  0],
       [14,  8,  6,  5,  7,  0]])

Solution 3:

A bit of indirect indexing does the trick. Pleaese note that I worked on rows since you started off on rows.

fdim = np.arange(3)[:, None]
so = np.argsort(myBigArray[fdim, top], axis=-1)
tops = top[fdim, so]
myBigArray[fdim, tops]
# array([[0, 1, 2],
         [0, 5, 6],
         [0, 5, 7]])

A note on argpartition with range argument: I strongly suspect that it is not O(n + k log k); in any case it is typically several-fold slower than a manual argpartition + argsort see here