How to use numpy.argsort() as indices in more than 2 dimensions? : <p>I know something similar to this question has been asked many times over already, but all answers given to similar questions only seem to work for arrays with 2 dimensions.</p>

<p>My understanding of <code>np.argsort()</code> is that <code>np.sort(array) == array[np.argsort(array)]</code> should be <code>True</code>.
I have found out that this is indeed correct if <code>np.ndim(array) == 2</code>, but it gives different results if <code>np.ndim(array) &gt; 2</code>.</p>

<p>Example:</p>

<pre><code>&gt;&gt;&gt; array = np.array([[[ 0.81774634,  0.62078744],
                       [ 0.43912609,  0.29718462]],
                      [[ 0.1266578 ,  0.82282054],
                       [ 0.98180375,  0.79134389]]])
&gt;&gt;&gt; np.sort(array)
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])
&gt;&gt;&gt; array.argsort()
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]])
&gt;&gt;&gt; array[array.argsort()]
array([[[[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]],



       [[[[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]],

         [[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]]])
</code></pre>

<p>So, can anybody explain to me how exactly <code>np.argsort()</code> can be used as the indices to obtain the sorted array?
The only way I can come up with is:</p>

<pre><code>args = np.argsort(array)
array_sort = np.zeros_like(array)
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        array_sort[i, j] = array[i, j, args[i, j]]
</code></pre>

<p>which is extremely tedious and cannot be generalized for any given number of dimensions.</p>
