寻找类似numpy.take的可赋值函数，支持动态指定操作轴

你可以用numpy.put_along_axis函数实现需求，它完美对应numpy.take的赋值场景，支持动态指定操作轴，不用写分支处理不同轴的情况，且能直接修改原数组。

基础示例（二维数组） #

比如你提到的二维数组：

import numpy as np
a = np.array([[1,2],[3,4]])

对应花式索引赋值a[0, :] = 10，用put_along_axis的写法是：

# 指定要修改的位置索引，需扩展维度匹配原数组形状
indices = np.array([0])[:, np.newaxis]
np.put_along_axis(a, indices, 10, axis=0)

执行后a会变为array([[10, 10], [3, 4]])，和花式索引的效果完全一致。

参数说明 #

• 第一个参数：要修改的目标数组
• 第二个参数：要操作的位置索引，注意索引的形状需要和目标数组在指定轴上的维度匹配（比如轴0是行维度，需要把一维索引转成列向量，用[:, np.newaxis]或[:, None]扩展维度）
• 第三个参数：要赋的值，可以是标量（会自动广播），也可以是和索引形状匹配的数组
• 第四个参数：指定要操作的轴

高维数组示例 #

对于三维数组，同样无需修改逻辑，直接指定轴即可：

b = np.arange(27).reshape(3,3,3)
# 修改轴1的第1个位置为0
indices = np.array([1])[None, :, None]
np.put_along_axis(b, indices, 0, axis=1)

另外，numpy.take_along_axis是numpy.take的对应取值函数，和put_along_axis配对使用，完全覆盖你需要的“非花式索引+动态轴+取值/赋值”场景。

内容的提问来源于stack exchange，提问作者5xum