基于时间索引数组与区域数组对三维数据数组进行聚合求和的实现方案咨询

看起来你需要根据时间分组索引和空间区域标签，对三维数据数组进行分组求和，最终得到每个时间组对应各个空间区域的总和。我这里提供两种高效的实现方案，分别基于NumPy（适合大规模数组运算）和Pandas（更直观易读），你可以根据自己的需求选择。

方案一：使用NumPy高效分组求和 #

NumPy的广播机制和np.add.at函数可以高效处理这种多维数组的分组聚合，避免循环带来的性能损耗。

import numpy as np

# 初始化你的原始数据
zon_arr = np.zeros((3,5))
tim_idx = np.array([0,0,1,1,2,2,3,3])
zon_arr[1, :3] = 1
zon_arr[1, 3:] = 2
np.random.seed(100)
dat_arr = np.random.rand(8, 3, 5)

# 第一步：获取唯一的时间标签和区域标签，确定结果的维度
unique_tim = np.unique(tim_idx)
unique_zones = np.unique(zon_arr)

# 初始化结果数组，行对应时间组，列对应区域
result = np.zeros((len(unique_tim), len(unique_zones)))

# 第二步：将时间索引和区域数组广播到与dat_arr相同的形状
# tim_idx扩展为(8,1,1)，自动广播到(8,3,5)
tim_broadcast = tim_idx[:, np.newaxis, np.newaxis]
# zon_arr扩展为(1,3,5)，自动广播到(8,3,5)
zone_broadcast = zon_arr[np.newaxis, :, :]

# 第三步：将原始标签转换为结果数组中的索引位置
tim_group_idx = np.searchsorted(unique_tim, tim_broadcast)
zone_group_idx = np.searchsorted(unique_zones, zone_broadcast)

# 第四步：使用np.add.at完成分组累加，这是NumPy中高效的分组求和方式
np.add.at(result, (tim_group_idx, zone_group_idx), dat_arr)

# 输出结果
print("NumPy聚合结果（行：时间组，列：区域）：")
print(result)
print("\n对应的时间组：", unique_tim)
print("对应的区域：", unique_zones)

方案解释： #

• 广播操作让每个数据点都能匹配到对应的时间组和区域标签，无需手动循环遍历；
• np.searchsorted用来将原始的时间/区域标签映射到结果数组的索引位置，确保每个分组对应结果的正确位置；
• np.add.at是专门针对多位置累加的函数，比普通循环快得多，尤其适合处理大尺寸数组。

方案二：使用Pandas直观分组求和 #

如果你更习惯用数据框的方式处理数据，Pandas的分组功能会更直观，代码可读性更强。

import pandas as pd
import numpy as np

# 初始化你的原始数据
zon_arr = np.zeros((3,5))
tim_idx = np.array([0,0,1,1,2,2,3,3])
zon_arr[1, :3] = 1
zon_arr[1, 3:] = 2
np.random.seed(100)
dat_arr = np.random.rand(8, 3, 5)

# 第一步：将三维数据展开为一维，同时生成对应的时间和区域标签
# np.repeat：将tim_idx的每个元素重复zon_arr.size次（对应每个空间点）
# np.tile：将zon_arr展平后重复len(tim_idx)次（对应每个时间点）
df = pd.DataFrame({
    'time': np.repeat(tim_idx, zon_arr.size),
    'zone': np.tile(zon_arr.flatten(), len(tim_idx)),
    'value': dat_arr.flatten()
})

# 第二步：按时间和区域分组，对数值求和
# unstack()将区域转换为列，得到二维结果（行=时间，列=区域）
result_df = df.groupby(['time', 'zone'])['value'].sum().unstack()

# 输出结果
print("Pandas聚合结果：")
print(result_df)

方案解释： #

• 通过np.repeat和np.tile将多维数组的维度信息转换为数据框的列，让每个数据点都有对应的时间和区域标签；
• groupby+sum的组合是Pandas标准的分组聚合操作，逻辑清晰，容易理解；
• unstack操作把区域标签从行索引转换为列，最终得到和NumPy方案一致的二维结果格式。

两种方案得到的结果完全一致，都能满足你的需求——每个时间组对应3个区域的求和值，总共4×3=12个结果（若某个时间组无对应区域的数据，会显示0或NaN，取决于方案）。

备注：内容来源于stack exchange，提问作者user2256085