numpy.take_along_axis#
- numpy.take_along_axis(arr, indices, axis)[source]#
通过匹配一维索引和数据切片从输入数组中获取值。
它沿着索引和数据数组中指定的轴迭代匹配的一维切片,并使用前者在后者中查找值。这些切片可以具有不同的长度。
返回沿轴的索引的函数,例如
argsort
和argpartition
,会为此函数生成合适的索引。版本 1.15.0 中的新功能。
- 参数:
- 返回值:
- out: ndarray (Ni…, J, Nk…)
索引结果。
另请参阅
take
沿着轴获取,对每个一维切片使用相同的索引
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]
示例
>>> import numpy as np
对于此示例数组
>>> 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]])
如果使用
keepdims
保留微不足道的维度,则同样的方法适用于最大值和最小值>>> np.max(a, axis=1, keepdims=True) array([[30], [60]]) >>> ai = np.argmax(a, axis=1, keepdims=True) >>> ai array([[1], [0]]) >>> np.take_along_axis(a, ai, axis=1) array([[30], [60]])
如果我们想同时获取最大值和最小值,我们可以先堆叠索引
>>> ai_min = np.argmin(a, axis=1, keepdims=True) >>> ai_max = np.argmax(a, axis=1, keepdims=True) >>> 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]])