Numpy 对每列的给定行进行切片

Question

我有一个 ... x n x m 数组，例如

a

，其中 ... 代表任意数量的附加维度。为了简单起见，我们将 n 维“行”和 m 维“列”称为“列”，即使数组是更高维的。

我还有一个长度为 n 的向量

v

，其中包含最后一个维度的索引（范围从 0 到 m-1）。我想创建一个数组

b

，它使用此向量来提取每行的指示列。

使用循环可以轻松地做到这一点。这是一个最小的工作示例：

import numpy as np

a = np.round(np.random.rand(2,3,4)*10)
v = [0, 2, 1]

print(a)
"""[[[ 1.  6.  9.  9.]
  [ 1.  8.  4. 10.]
  [ 0.  0.  5.  3.]]

 [[ 7.  8.  1. 10.]
  [ 7.  9.  7.  8.]
  [ 3.  4.  8.  7.]]]
"""

b = []
for i in range(len(v)):
    b.append(a.take(i, axis=-2).take(v[i], axis=-1))
b = np.asarray(b)

print(b)
"""
[[1. 7.]
 [4. 7.]
 [0. 4.]]
"""

有没有更聪明的方法可以在不循环的情况下进行这种索引？

Answer 1

我想这会成功：

a[:, np.arange(len(v)), v].T

一般来说，您可以使用此语法对所有维度进行索引。在上面的示例中，您使用

:

选择第一个维度（深度）中的所有元素，然后使用

np.arange(len(v))

选择第二个维度（行）中的元素 0, 1, 2，最后选择元素 0, 2,最后一个维度（列）为 1，带有

v

。我将结果与

T

转置，以获得与

b

中相同的尺寸。

来自文档：

切片元组中具有多个非：条目的基本切片，其作用类似于使用单个非：条目重复应用切片，其中连续获取非：条目（所有其他非：条目替换为： ）。因此，x[ind1, ..., ind2,:] 的作用类似于基本切片下的 x[ind1][..., ind2, :]。

Answer 2

这个怎么样，假设

i

是 (0,1,2) 且

v

是 (0,2,1):

ind = [slice(None)] * (a.ndim - 2) + [np.array(i), np.array(v)]
k = a[ind]
transpose_dim = [a.ndim-2]+  list(range(0, a.ndim-2)) 
k = k.transpose(transpose_dim)

这个

[slice(None)] * (a.ndim - 2)

得到

a[:,:,.....]

，直到

a

的最后两个变暗，并且

[np.array(i), np.array(v)]

从每一行中选择正确的索引。然后你必须将

k

的行转换为

transpose

的列。