NumPy参考 >例行程序 >Indexing routines > numpy.take_along_axis
numpy.
take_along_axis
(arr,indexs,axis )[源代码] ¶通过匹配1d索引和数据切片从输入数组中获取值。
这将迭代沿索引和数据数组中沿指定轴定向的匹配1d切片,并使用前者在后者中查找值。这些切片的长度可以不同。
沿轴返回索引的函数(例如argsort
和
argpartition
)会为此函数生成合适的索引。
1.15.0版中的新功能。
索引结果。
也可以看看
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]])