如何利用Numpy切片替代循环实现四维数组批量赋值？

直接用Numpy的广播机制就能搞定，完全不需要循环，代码简洁又高效：

方法1：新增维度实现自动填充 #

import numpy as np

# 假设num_s是你的维度大小，values是(num_s, num_s)的数组
pair_values = values[:, :, np.newaxis, np.newaxis]

等价写法（任选其一） #

• 用reshape明确指定维度：

  pair_values = values.reshape(num_s, num_s, 1, 1)
  

• 用expand_dims指定新增维度的位置：

  pair_values = np.expand_dims(values, axis=(2, 3))
  

原理说明 #

原来的values是(num_s, num_s)的二维数组，我们给它新增两个长度为1的维度，变成(num_s, num_s, 1, 1)的四维数组。根据Numpy的广播规则，当把这个数组赋值给(num_s, num_s, num_s, num_s)的目标数组时，长度为1的维度会自动扩展成num_s，相当于把每个values[x,y]的标量值自动填充到对应的num_s×num_s子数组里，和你原来的循环效果完全一致，但效率高得多（尤其是num_s很大时）。

验证示例 #

比如num_s=2，values = np.array([[1,2],[3,4]])，用上面的方法得到的pair_values：

• pair_values[0,0]会是[[1,1],[1,1]]
• pair_values[0,1]会是[[2,2],[2,2]]
• pair_values[1,0]会是[[3,3],[3,3]]
• pair_values[1,1]会是[[4,4],[4,4]]

完全符合你的需求，而且没有任何循环。

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