Python-Numpy库的扩维(维度扩展)操作示例代码积累

Numpy库中有不少与索引相关的方法或函数要求索引的维度与输入数组的维度一致，比如函数take_along_axis()和函数put_along_axis()都要求这样。

所以本篇博文记录下扩维(维度扩展)的操作示例代码，以便自己和大家需要时方便取用。

第01种方法：使用函数expand_dims()

import numpy as np

A = np.array([[10, 30, 20], [60, 40, 50]])
A_max_index = np.argmax(A, axis=1)

A_max_index_expand = np.expand_dims(A_max_index, axis=1)

运行结果如下：

第02种方法：使用None对象
示例代码如下：

import numpy as np

A = np.array([[10, 30, 20], [60, 40, 50]])
A_max_index = np.argmax(A, axis=1)

A_max_index_expand = A_max_index[:, None]

运行结果如下：