如何用NumPy/Pandas对数组非零簇求和并将总和置于簇中间

针对你给出的需求，这里提供两种高效实现方法，分别基于NumPy和Pandas：

NumPy实现（推荐，适合大型数组） #

import numpy as np

# 示例数组
test = np.array([0, 0, 1, 3, 5, 0, 0, 0, 15, 16, 2, 0, 0])

# 标记非零元素
non_zero = test != 0
# 检测簇的起始和结束位置
changes = np.diff(np.concatenate([[0], non_zero, [0]]))
starts = np.where(changes == 1)[0]
ends = np.where(changes == -1)[0] - 1

# 计算每个簇的和与中间索引
cluster_sums = [test[s:e+1].sum() for s, e in zip(starts, ends)]
middle_indices = [(s + e) // 2 for s, e in zip(starts, ends)]

# 生成结果数组
result = np.zeros_like(test)
for idx, val in zip(middle_indices, cluster_sums):
    result[idx] = val

print(result)
# 输出：[ 0  0  0  9  0  0  0  0  0 33  0  0  0]

逻辑说明 #

1. 先标记数组中的非零元素，通过np.diff检测非零状态的变化，定位每个连续非零簇的起始和结束索引。
2. 对每个簇求和，并计算簇的中间位置索引（奇数长度取正中间，偶数长度取靠左的中间位置，若需要靠右可调整为(s + e + 1)//2）。
3. 创建全零数组，将每个簇的和填入对应的中间位置。

Pandas实现（直观易读） #

如果你更熟悉Pandas，也可以用分组的方式实现：

import pandas as pd
import numpy as np

test = np.array([0, 0, 1, 3, 5, 0, 0, 0, 15, 16, 2, 0, 0])
df = pd.DataFrame({'val': test})

# 生成簇的分组键：连续非零元素归为同一组
df['cluster'] = (df['val'] == 0).cumsum()
# 筛选出非零簇并分组
non_zero_clusters = df[df['val'] != 0].groupby('cluster')

# 生成结果数组
result_pd = np.zeros_like(test)
for _, group in non_zero_clusters:
    cluster_sum = group['val'].sum()
    # 计算中间索引，和NumPy方法逻辑一致
    middle_idx = (group.index[0] + group.index[-1]) // 2
    result_pd[middle_idx] = cluster_sum

print(result_pd)
# 输出：[ 0  0  0  9  0  0  0  0  0 33  0  0  0]

这两种方法都可以直接适配你提供的真实数据集（包含小数和0的长数组），只需要把示例数组替换成你的真实数据即可。

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