numpy根据索引列表选择值。一次处理批次

Question

一般任务似乎很简单。 我给了两个数组： 数组A of shape (5,20)，即五个100个值的列表。 另一个数组B of shape (5, 10)包含我要选择的索引。 [因此数组B包含0到19之间的值]

一个最小的例子将由

解决

# saqmple initialization
def rnd():
    return np.random.choice(20, 10, replace=False)
A = np.random.random_sample((5, 20))
B = np.array([rnd(),  rnd(),  rnd(),  rnd(),  rnd()])

# solution
selected_values = [A[i,B[i]] for i in range(A.shape[0])]

以更简单的示例为例

A = [[4,5,6],[7,8,9]] 
B = [[0,1],[1,2]]
out = [[4,5],[8,9]]

有没有一种方法可以重塑，重新格式化或...从而能够在没有for循环的情况下执行此操作？即隐含地让numpy意识到B中的第二个条目是对应于A中的第二个条目的索引，依此类推？