技术背景
在前一篇文章中，我们提到了关于Numpy中的各种取index的方法，可以用于取出数组里面的元素，也可以用于做切片，甚至可以用来做排序。但是遇到对于高维矩阵的某一个维度取多个值的时候，单纯的使用下标已经无法完成相关的操作了。如果找不到相应的接口，对于性能要求不高的场景可以使用一个for循环进行替代，但是对于性能要求比较高的场景下，我们还是尽可能的使用Numpy本身自带的接口，比如本文将要提到的take_along_axis操作。
使用案例
我们考虑这样的一个场景，给定一个维度为(4,11,3)的矩阵a作为数据，和一个维度为(4,2)的矩阵b作为下标，意味着从a中第二条轴的11个元素中每次取两个元素，也就是希望得到一个维度为(4,2,3)的结果：
In [11]: a = np.arange(132).reshape((4,11,3))
In [12]: a
Out[12]:
array([[[ 0, 1, 2],
[ 3, 4, 5],
[ 6, 7, 8],
[ 9, 10, 11],
[ 12, 13, 14],
[ 15, 16, 17],
[ 18, 19, 20],
[ 21, 22, 23],
[ 24, 25, 26],
[ 27, 28, 29],
[ 30, 31, 32]],
[[ 33, 34, 35],
[ 36, 37, 38],
[ 39, 40, 41],
[ 42, 43, 44],
[ 45, 46, 47],
[ 48, 49, 50],
[ 51, 52, 53],
[ 54, 55, 56],
[ 57, 58, 59],
[ 60, 61, 62],
[ 63, 64, 65]],
[[ 66, 67, 68],
[ 69, 70, 71],
[ 72, 73, 74],
[ 75, 76, 77],
[ 78, 79, 80],
[ 81, 82, 83],
[ 84, 85, 86],
[ 87, 88, 89],
[ 90, 91, 92],
[ 93, 94, 95],
[ 96, 97, 98]],
[[ 99, 100, 101],
[102, 103, 104],
[105, 106, 107],
[108, 109, 110],
[111, 112, 113],
[114, 115, 116],
[117, 118, 119],
[120, 121, 122],
[123, 124, 125],
[126, 127, 128],
[129, 130, 131]]])
In [13]: b = np.array([[0,1],[1,2],[2,3],[3,4]])
In [14]: b
Out[14]:
array([[0, 1],
[1, 2],
[2, 3],
[3, 4]])
为了方便展示我们就定义了这样两个比较简单的矩阵a和b，那么在这个结果中，我们理想的结果应该是：
[[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 36, 37, 38],
[ 39, 40, 41]],
[[ 72, 73, 74],
[ 75, 76, 77]],
[[108, 109, 110],
[111, 112, 113]]]
这样的一个矩阵。关于这个结果的来源，可以对b这个定义进行展开解释，b的值为：
[[0, 1],
[1, 2],
[2, 3],
[3, 4]]
它所表示的是在a[0]下取第0个元素和第1个元素，在a[1]下取第1个元素和第2个元素，以此类推。然而如果我们直接把定义好的b放到a的索引中或者直接使用numpy.take的方法的话，得到的结果是这样的：
In [16]: a[:,b]
Out[16]:
array([[[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 3, 4, 5],
[ 6, 7, 8]],
[[ 6, 7, 8],
[ 9, 10, 11]],
[[ 9, 10, 11],
[ 12, 13, 14]]],
[[[ 33, 34, 35],
[ 36, 37, 38]],
[[ 36, 37, 38],
[ 39, 40, 41]],
[[ 39, 40, 41],
[ 42, 43, 44]],
[[ 42, 43, 44],
[ 45, 46, 47]]],
[[[ 66, 67, 68],
[ 69, 70, 71]],
[[ 69, 70, 71],
[ 72, 73, 74]],
[[ 72, 73, 74],
[ 75, 76, 77]],
[[ 75, 76, 77],
[ 78, 79, 80]]],
[[[ 99, 100, 101],
[102, 103, 104]],
[[102, 103, 104],
[105, 106, 107]],
[[105, 106, 107],
[108, 109, 110]],
[[108, 109, 110],
[111, 112, 113]]]])
显然这不是我们想要的结果。需要额外申明的是，这个执行操作中，最后一个维度的冒号加与不加是一样的效果，跟numpy.take本质上也是同样的操作，因此就需要使用到numpy中的另外一个接口：take_along_axis，如下是其官方的API文档：
还有相关的使用案例：
需要注意的是，输入的indices必须要跟原始的数据矩阵保持同样的维度，因此在我们自己的案例中，对b进行了扩维，最终的代码如下所示：
In [23]: np.take_along_axis(a,b[:,:,None],axis=1)
Out[23]:
array([[[ 0, 1, 2],
[ 3, 4, 5]],
[[ 36, 37, 38],
[ 39, 40, 41]],
[[ 72, 73, 74],
[ 75, 76, 77]],
[[108, 109, 110],
[111, 112, 113]]])
最后得到的就是我们想要的结果了，并且是直接使用下标无法实现的操作（当然，也可能是我还没研究出来这样的操作）。这里axis设置为1，就表示a的第0个维度和b的第0个维度是一致的取法，也可以理解成全取的意思。
总结概要
Numpy是在Python中用于各种矩阵运算非常强大的工具之一，而快速的通过下标取出所需位置的元素也是numpy所支持的强大功能之一。常规的元素取法都可以通过numpy的下标或者是numpy.take函数来实现，比如array[0,:]可用于取第一条轴的所有元素，array[:,0]可以用于取第二条轴的所有第二个元素，放在一个2维的矩阵里面就分别是取第一行的所有元素和取第一列的所有元素。但是本文更加关注于更高维的矩阵，当我们想从多个维度中取多个元素时，是不太容易直接用下标去取的，比如同时取a[0][0]，a[0][1]，a[1,1]，a[1][2]的话，那么就只能使用numpy所支持的另外一个函数numpy.take_along_axis来实现。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/take_along_axis.html
作者ID：DechinPhy
更多原著文章请参考：https://www.cnblogs.com/dechinphy/
打赏专用链接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
腾讯云专栏同步：https://cloud.tencent.com/developer/column/91958