def original_app(x, idx):
    _, yt, xt = x.shape
y = np.zeros((yt, xt))
for j in range(yt):
    for i in range(xt):
        y[j, i] = x[idx[j, i], j, i]
return y

def ogrid_based(x, idx):
    _, yt, xt = x.shape
J, I = np.ogrid[:yt, :xt]
return x[idx, J, I]

def reshape_based(x, idx):
    zt, yt, xt = x.shape
return x.reshape(zt, -1)[idx.ravel(), np.arange(yt * xt)].reshape(-1, xt)
NumPy: оценка массива индексов во время векторизованного назначения