基于连续相同值拆分Numpy数组后，求高效数值调整优化方案

问题概述 #

给定如下Numpy数组：

import numpy as np
arr = np.array([1,1,1,2,2,2,3,3,2,2,2,1,1,1,2,2])

我们需要先将其拆分为连续相同值的子数组，然后对每个子数组执行规则：若子数组的唯一元素此前已出现过，则为该子数组的每个元素加上 0.001 * 该元素此前出现的次数（原代码中使用的是0.0001，可根据需求调整）。

以下是比原实现更高效的向量化方案，避免了多次循环和数组拼接操作。

优化代码实现 #

import numpy as np

arr = np.array([1,1,1,2,2,2,3,3,2,2,2,1,1,1,2,2])

# 标记连续子数组的边界，生成每个元素对应的组ID
diff = np.diff(arr, prepend=arr[0] - 1)
group_ids = np.cumsum(diff != 0)

# 获取每个组的唯一值
group_vals = arr[diff != 0]

# 统计每个值的组出现次数（按顺序）
counts = np.zeros(len(group_vals), dtype=int)
for val in np.unique(group_vals):
    # 找到当前值对应的所有组位置
    val_indices = np.where(group_vals == val)[0]
    # 为这些组分配递增的计数（从0开始）
    counts[val_indices] = np.arange(len(val_indices))

# 将计数映射到原数组的每个元素，计算增量
increments = counts[group_ids - 1] * 0.001  # 若需原代码的0.0001，替换此处即可
result = arr.astype(float) + increments

print(result)

输出结果 #

按需求中的0.001计算，输出为：

array([1.    , 1.    , 1.    , 2.    , 2.    , 2.    , 3.    , 3.    ,
       2.001 , 2.001 , 2.001 , 1.001 , 1.001 , 1.001 , 2.002 , 2.002 ])

若替换为0.0001，则与原代码结果完全一致。

方案优势 #

1. 避免数组拆分：通过生成组ID的方式，无需实际拆分数组，减少内存开销和冗余操作步骤。
2. 向量化为主：核心逻辑仅用一次循环处理不同值的计数，其余操作均为Numpy向量化操作，比原实现的多次循环效率更高。
3. 逻辑直观清晰：从标记组边界到统计次数，再到计算增量，每一步逻辑明确，便于后续维护和调整规则。

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