title: sorting along multidimensional index in numpy
body: <p>I have a three dimensional array, that I need to sort. I also have a 2 dimensional index that I got with a <code>lexsort</code> that I would like to use for sorting. I cannot find how to apply the index to the array without extra dimensions sprouting up...</p>
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  
  a = np.array([
      [
          [1, 1, 1],
          [0, 1, 2],
          [3, 1, 1],
          [-1, 1, 1],
          [0, 2, 2],
          [1, 2, 2],
      ],[
          [-1, 1, 1],
          [0, 1, 2],
          [-3, 1, 1],
          [1, 1, 1],
          [0, -2, 2],
          [-1, -2, 2],
      ],[
          [-1, 1, 1],
          [-0, 1, -2],
          [-3, 1, 1],
          [1, 1, 1],
          [-0, -2, -2],
          [-1, -2, -2],
      ],
  ])
  sorted_index = np.lexsort((a[:,:,0], a[:,:,1],a[:,:,2]),axis=1)
  
  </code></pre>
  <p>sorted_index is</p>
  <pre class="lang-py prettyprint-override"><code>[[3 0 2 1 4 5], [2 0 3 5 4 1], [5 4 1 2 0 3]]
  </code></pre>
  <p>the end result should be</p>
  <pre class="lang-py prettyprint-override"><code>a_sorted = np.array([
      [
          [-1, 1, 1],
          [1, 1, 1],
          [3, 1, 1],
          [0, 1, 2],
          [0, 2, 2],
          [1, 2, 2],
      ],[
          [-3, 1, 1],
          [-1, 1, 1],
          [1, 1, 1],
          [-1, -2, 2],
          [0, -2, 2],
          [0, 1, 2],
      ],[
          [-1, -2, -2],
          [0, -2, -2],
          [0, 1, -2],  
          [-3, 1, 1],
          [-1, 1, 1],
          [1, 1, 1],
      ],
  ])
  
  </code></pre>
  <p>I tried.</p>
  <pre class="lang-py prettyprint-override"><code>a[sorted_index]
  a[sorted_index, :]
  a[:, sorted_index]
  </code></pre>
  <p>I am about to loop through the thing, but I would like to do it properly.</p>