矩阵上的快速numpy行切片

Question

我有以下问题：我有一个大小为(m,200)（m = 3683）的矩阵yj，我有一个字典，它为每个键返回一个行索引的numpy数组yj（对于每个键，大小数组都会改变，以防万一有人想知道）

现在，我必须多次访问这个矩阵（大约100万次），我的代码由于索引而变慢了（我已经分析了代码，这一步花费了65%的时间）

以下是我尝试过的：

1. 首先，使用索引进行切片：

>> %timeit yj[R_u_idx_train[1]]
10.5 µs ± 79.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

变量R_u_idx_train是具有行索引的字典

1. 我认为布尔索引可能更快：

>> yj[R_u_idx_train_mask[1]]
10.5 µs ± 159 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

R_u_idx_train_mask是一个字典，它返回一个大小为m的布尔数组，其中R_u_idx_train给出的索引设置为True

1. 我也试过np.ix_

>> cols = np.arange(0,200)
>> %timeit ix_ = np.ix_(R_u_idx_train[1], cols); yj[ix_]
42.1 µs ± 353 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

1. 我也试过np.take

>> %timeit np.take(yj, R_u_idx_train[1], axis=0)
2.35 ms ± 88.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

虽然这看起来很好，但事实并非如此，因为它给出的数组是shape(R_u_idx_train[1].shape[0], R_u_idx_train[1].shape[0])（应该是(R_u_idx_train[1].shape[0], 200)）。我想我没有正确使用这个方法

1. 我也试过np.compress

>> %timeit np.compress(R_u_idx_train_mask[1], yj, axis=0)
14.1 µs ± 124 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

1. 最后，我尝试用布尔矩阵索引

>> %timeit yj[R_u_idx_train_mask2[1]]
244 µs ± 786 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

那么10.5 µs ± 79.7 ns per loop是我能做的最好的了吗？我可以尝试使用cython，但这似乎只是索引的大量工作

非常感谢

Answer 1

V.Ayrat在评论中给出了一个非常聪明的解决方案

>> newdict = {k: yj[R_u_idx_train[k]] for k in R_u_idx_train.keys()}
>> %timeit newdict[1]
202 ns ± 6.7 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)

无论如何，如果知道是否有一种方法可以使用numpy来加速它，也许还是很酷的