如何修改任意维度多维数组的最后一维元素？

问题场景 #

对于二维数组：

import numpy as np
test_array = np.ones((2,2))
test_array[:,1] = -99  # 修改最后一维的第2列

对于三维数组：

test_array = np.ones((2,2,2))
test_array[:,:,1] = -99  # 修改最后一维的第2层

现在需要编写通用代码，处理任意n维数组（比如test_array = np.ones([2]*n)），实现对最后一维指定位置的赋值，本质是要自动生成对应数量的:来匹配数组维度。

解决方案：用...（Ellipsis）简化索引 #

在numpy的索引规则里，...代表“匹配剩余所有维度的全部元素”，刚好可以用来替代任意数量的:，直接实现任意维度数组的最后一维赋值：

# 任意n维数组都适用
test_array[..., 1] = -99

示例验证 #

• 4维数组：

  test_array = np.ones((2,2,2,2))
  test_array[..., 1] = -99
  

  这行代码等价于test_array[:,:,:,1] = -99，会把最后一维索引为1的所有元素设为-99。

• 封装成函数：
  如果需要复用逻辑，可以写成函数：

  def set_last_dimension(arr, target_value, index=1):
      arr[..., index] = target_value
  

  调用时直接传入数组、目标值和要修改的最后一维索引即可，无需关心数组维度。

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