numpy.take_along_axis

numpy.take_along_axis(arr, indices, axis)

通过匹配一维索引和数据切片从输入数组中获取值。

这将遍历索引和数据数组中沿指定轴定向的匹配1d切片，并使用前者查找后者中的值。这些切片可以是不同的长度。

函数沿轴返回索引，如 argsort 和 argpartition ，为此函数生成适当的索引。

1.15.0 新版功能.

参数

阿瑞：恩达瑞（ni…，m，nk…）

源阵列

索引：ndarray（ni…，j，nk…）

每一维切片的索引 arr . 这必须与ARR的尺寸相匹配，但NI和NJ的尺寸只需广播 arr .

轴：int

要沿其进行1d切片的轴。如果“轴”为“无”，则将输入数组视为其第一次被展平为1d，以便与 sort 和 argsort .

返回

输出：Ndarray（ni…，j，nk…）

索引结果。

参见

take

沿着一个轴，对每个1d切片使用相同的索引

put_along_axis

通过匹配一维索引和数据切片将值放入目标数组

笔记

这相当于（但比）以下使用 ndindex 和 s_ ，设置每个 ii 和 kk 到一组索引：

Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
J = indices.shape[axis]  # Need not equal M
out = np.empty(Ni + (J,) + Nk)

for ii in ndindex(Ni):
    for kk in ndindex(Nk):
        a_1d       = a      [ii + s_[:,] + kk]
        indices_1d = indices[ii + s_[:,] + kk]
        out_1d     = out    [ii + s_[:,] + kk]
        for j in range(J):
            out_1d[j] = a_1d[indices_1d[j]]

同样地，消除内环，最后两行是：

out_1d[:] = a_1d[indices_1d]

实例

对于此示例数组

>>> a = np.array([[10, 30, 20], [60, 40, 50]])

我们可以直接使用sort或argsort和此函数进行排序

>>> np.sort(a, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])
>>> ai = np.argsort(a, axis=1); ai
array([[0, 2, 1],
       [1, 2, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 20, 30],
       [40, 50, 60]])

如果展开尺寸，最大值和最小值也适用：

>>> np.expand_dims(np.max(a, axis=1), axis=1)
array([[30],
       [60]])
>>> ai = np.expand_dims(np.argmax(a, axis=1), axis=1)
>>> ai
array([[1],
       [0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[30],
       [60]])

如果我们想同时得到最大值和最小值，我们可以先将索引堆叠起来。

>>> ai_min = np.expand_dims(np.argmin(a, axis=1), axis=1)
>>> ai_max = np.expand_dims(np.argmax(a, axis=1), axis=1)
>>> ai = np.concatenate([ai_min, ai_max], axis=1)
>>> ai
array([[0, 1],
       [1, 0]])
>>> np.take_along_axis(a, ai, axis=1)
array([[10, 30],
       [40, 60]])

numpy.take numpy.choose