python - 如何在不使用numpy或tensorflow中的任何循环的情况下获取矩阵行？

Question

我有一个矩阵，我想读取矩阵的每一行并使用 tf.nn.top_k 查找每行的前 k 值。

如何在不使用循环或列表理解的情况下获取矩阵的每一行。我想使用 Tensorflow 或 numpy，并且正在考虑应用广播。

如果我将索引放入数组中，例如矩阵中有 10 行，我会:

indices = [0,1,2,3......,9]

然后我可以应用广播吗？

Answer 1

说(就像这里的另一个答案一样):

a = np.random.randint(0, 1000, (4,4))

那么你可以简单地做:

np.sort(a)[:,-1:-3:-1]

其中而不是 3您在索引中输入了一个与您想要获取的最高值数量相关的数字。

编辑:为了解决您的评论:首先，我更改为 float 数据类型(以处理 1/3 值):

a = 1000*np.random.random((4,4))

然后我得到索引:

idx = a.argsort().argsort()

将 1/3 分配给前两个值:

a[idx>=2] = 1./3.

并将零分配给其余部分:

a[idx<2] = 0

瞧。