[FIXED] Find k smallest values of of a 2D array slice of a 3D Array?

Issue

I've read through other questions where other people are trying to find the k smallest values with np.partition or np.argpartition. However, these seem like simple 1D or 2D arrays to operate on.

What I have is a 3D array where I'm taking a 2D slice of it and trying to find the 4 smallest values as such:

import numpy as np
np.random.seed(556)

big_array = np.random.randint(0,200,size = (15,9,3))
for i in range(3):
    min_vals = np.sort(big_array[:,:,i].flatten())[:4]
    print("The smallest values for last axis #{} are: {:.1f},{:.1f},{:.1f},{:.1f}".format(i,*min_vals))

Unfortunately, it seems like np.partition doesn't like it when the kth argument is larger than the dimension of whatever axis you're trying to specify, probably because I don't fully comprehend the syntax. Likewise, it doesn't look like np.sort or np.argsort have options for multiple axes arguments.

Is there a way to accomplish this without using a for loop or list comprehension?

Solution

IIUC, you can sort the reshaped array along the first axis and extract the results from there. Your loop and flatten is the same as:

big_array.reshape(-1, big_array.shape[-1])

So, your code would be:

top_4_per_last_axis = np.sort(big_array.reshape(-1, big_array.shape[-1]), 0)[:4]

Result:

array([[ 0,  2,  0],
       [ 1,  3,  2],
       [ 1,  6,  3],
       [ 2,  9, 12]])

Answered By - jared