numpy :применение argsort к массиву

argsort()Функция возвращает матрицу индексов, которую можно использовать для индексации исходного массива, чтобы результат соответствовал результату sort().

Есть ли способ применить эти индексы? У меня есть два массива, один из которых используется для получения порядка сортировки, а другой — некоторые связанные данные.

Я хотел бы вычислить assoc_data[array1.argsort()], но это не работает.

Вот пример:

z=array([1,2,3,4,5,6,7])
z2=array([z,z*z-7])
i=z2.argsort()

z2=array([[ 1,  2,  3,  4,  5,  6,  7],
          [-6, -3,  2,  9, 18, 29, 42]])
i =array([[1, 1, 1, 0, 0, 0, 0],
          [0, 0, 0, 1, 1, 1, 1]])

Я хотел бы применить i к z2 (или другому массиву со связанными данными ), но не знаю, как это сделать.