如何用Numpy基于数组B对3D数组A下采样生成数组C

问题描述 #

现有两个Numpy数组：

• 数组A：形状为(20,5,5)，由20个5×5的二维数组堆叠而成
• 数组B：形状为(20,)，值为[0,8]范围内的递增整数（含重复），具体值为：

  array([0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 7, 7, 7, 8, 8, 8])
  

需要生成形状为(3,5,5)的数组C：

• C的第一维度对应B中的值下采样到目标集合[0,4,8]
• C的每个元素是A中对应B值分组的均值，例如C[2,0,0]等于A中所有B值为8的元素在(0,0)位置的均值

生成A、B及目标采样点的代码如下：

import numpy as np
z = 20
x = 5
y = 5

# Array A
A = np.arange(1, 501,1)
A = A.reshape(z, x, y)
print ("A.shape = {}".format(A.shape)) # (20, 5, 5)

# Array B
B = np.arange(0, 12, 1)
B = np.repeat(B, 3)
B = B[[2,3,4,6,7,8,9,10,12,13,15,16,17,20,21,22,23,24,25,26]]
print ("B.shape = {}".format(B.shape)) # (20,)
print ("Array B = {}".format(B)) # array([0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 5, 6, 7, 7, 7, 8, 8, 8])

# 目标采样点
t = np.arange(0, 12, 4)
print ("Sampling by 4 = {}".format(t)) # array([0, 4, 8])

解决方案 #

方法1：循环遍历目标采样点 #

直接针对每个目标采样点，筛选A中对应切片并计算均值，逻辑直观易懂：

# 初始化结果数组C
C = np.zeros((len(t), x, y))

for i, target_val in enumerate(t):
    # 生成布尔掩码，筛选B中等于目标值的位置
    mask = B == target_val
    # 取A中对应位置的元素，沿z轴（第一个维度）计算均值
    C[i] = A[mask].mean(axis=0)

# 验证结果
print("C.shape:", C.shape)  # 输出 (3,5,5)
print("C[2,0,0]:", C[2,0,0])  # 输出 476.0，对应A[17,0,0]=451、A[18,0,0]=476、A[19,0,0]=501的均值

方法2：预计算所有分组均值再提取目标值 #

先计算B中所有唯一值对应的分组均值，再提取目标采样点的结果，适合需要多次提取不同采样点的场景：

# 获取B中的唯一值
unique_b_vals = np.unique(B)
# 初始化所有分组的均值数组，形状为(max(B)+1,5,5)
all_group_means = np.zeros((unique_b_vals.max() + 1, x, y))

for b_val in unique_b_vals:
    # 计算当前B值对应的A切片的均值
    all_group_means[b_val] = A[B == b_val].mean(axis=0)

# 提取目标采样点对应的均值，得到数组C
C = all_group_means[t]

# 验证结果
print("C.shape:", C.shape)  # 输出 (3,5,5)
print("C[0,0,0]:", C[0,0,0])  # 输出 1.0，对应A[0,0,0]的均值（仅一个元素）

关键逻辑说明 #

• 利用布尔掩码筛选A中与B值匹配的元素：A[B == target_val]会得到形状为(n,5,5)的数组，其中n是B中等于target_val的元素个数
• 沿第一个维度（z轴）调用mean(axis=0)，对每个5×5位置的元素取均值，得到单个5×5数组，赋值给C的对应索引位置

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