Python 用形状相同的数组中的相应行逐行切片二维数组

考虑数组

a

np.random.seed([3,1415])
a = np.random.choice(np.arange(8), (2, 4), False)
print(a)

[[7 1 4 5]
 [6 2 3 0]]

我将创建另一个数组

b

，它沿每一行保存

np.argsort

的结果

b = a.argsort(1)
print(b)

[[1 2 3 0]
 [3 1 2 0]]

我想通过使用

b

对

a

进行适当的切片来生成排序后的

a

。应该是这样的

idx0 = np.arange(a.shape[0]).repeat(a.shape[1])
print(a[idx0, b.ravel()].reshape(a.shape))

[[1 4 5 7]
 [0 2 3 6]]

问题
以上述方式用另一个

2x4

数组切片

2x4

数组的合适方法是什么？

获取帮助-

a[np.arange(b.shape[0])[:,None],b]

样本运行-

In [10]: a
Out[10]: 
array([[7, 1, 4, 5],
       [6, 2, 3, 0]])

In [11]: b
Out[11]: 
array([[1, 2, 3, 0],
       [3, 1, 2, 0]])

In [12]: a[np.arange(b.shape[0])[:,None],b]
Out[12]: 
array([[1, 4, 5, 7],
       [0, 2, 3, 6]])