优化NumPy代码：连续数值替换提速及数组关联问题咨询

原代码与问题 #

用户提供的原代码如下：

import numpy as np

# an array to split
a = np.array([0,0,1,0,1,1,1,0,1,1,0,0,0,1])

# idx where the number changes
idx = np.where(np.roll(a,1)!=a)[0][1:]

# split of array into groups
aout = np.split(a,idx)

# sum of each group
sumseg = [aa.sum() for aa in aout]

#fill criteria
idx2 = np.where( (np.array(sumseg)>0) & (np.array(sumseg)<2) )

#fill targets
[aout[ai].fill(0) for ai in idx2[0]]

# a is now updated? didn't follow how a gets updated
# return a

用户的疑问：为什么拆分后的aout修改会影响原数组a？实际场景中需要对二维数组的每行/列执行该操作，希望优化运行速度。

为什么修改aout会更新原数组a #

np.split()函数返回的是原数组的视图（view），而非独立的副本（copy）。视图与原数组共享同一块内存空间，修改aout中的子数组时，本质是直接修改原数组a对应的内存区域，因此原数组会同步更新。若要避免这种关联，可在拆分时对每个子数组调用.copy()，比如aout = [arr.copy() for arr in np.split(a, idx)]，但这会额外占用内存并降低效率。

代码优化方案 #

原代码多次使用列表推导和循环，对大数组或二维数组效率极低。以下是纯NumPy向量化的优化版本，彻底避免循环，大幅提升运行速度：

1. 一维数组优化版本 #

import numpy as np

a = np.array([0,0,1,0,1,1,1,0,1,1,0,0,0,1])

# 找到数组元素变化的位置，替代原代码的idx
# np.diff(a) != 0 标记相邻元素不同的位置，+1得到split的分割索引
split_idx = np.where(np.diff(a) != 0)[0] + 1
# 补充首尾索引，方便用reduceat计算分段求和
reduce_idx = np.concatenate([[0], split_idx, [len(a)]])

# 用np.add.reduceat替代循环求和，效率远高于列表推导
sumseg = np.add.reduceat(a, reduce_idx[:-1])

# 定位需要填充0的分段：sum>0且sum<2（即sum=1）
mask = (sumseg > 0) & (sumseg < 2)

# 生成需要置0的位置掩码
fill_mask = np.zeros(len(a), dtype=bool)
for start, end, valid in zip(reduce_idx[:-1], reduce_idx[1:], mask):
    if valid:
        fill_mask[start:end] = True

# 直接修改原数组
a[fill_mask] = 0

2. 二维数组按行处理的优化版本 #

针对二维数组的行/列处理，提供两种方案：

方案A：使用np.apply_along_axis（简洁易读）

import numpy as np

def process_row(row):
    split_idx = np.where(np.diff(row) != 0)[0] + 1
    reduce_idx = np.concatenate([[0], split_idx, [len(row)]])
    sumseg = np.add.reduceat(row, reduce_idx[:-1])
    mask = (sumseg > 0) & (sumseg < 2)
    fill_mask = np.zeros(len(row), dtype=bool)
    for s, e, v in zip(reduce_idx[:-1], reduce_idx[1:], mask):
        if v:
            fill_mask[s:e] = True
    row[fill_mask] = 0
    return row

# 示例二维数组
a_2d = np.array([
    [0,0,1,0,1,1],
    [1,0,0,1,0,1],
    [0,1,0,0,1,1]
])

# 按行处理
processed_2d = np.apply_along_axis(process_row, axis=1, arr=a_2d)

方案B：完全向量化（适合超大数组，效率最高）

完全向量化避免Python层循环，充分利用NumPy的C级运算效率：

import numpy as np

a_2d = np.array([
    [0,0,1,0,1,1],
    [1,0,0,1,0,1],
    [0,1,0,0,1,1]
])

rows, cols = a_2d.shape

# 计算每行的相邻元素差异，标记变化点
diff = np.diff(a_2d, axis=1) != 0
# 为每行添加首尾的0（标记行首和行尾不是变化点）
diff_padded = np.pad(diff, ((0,0), (1,1)), mode='constant')
# 计算每行的分段ID（累积求和标记不同分段）
segment_ids = np.cumsum(diff_padded, axis=1)

# 按行计算每个分段的和
sumseg = np.array([np.bincount(seg, weights=row) for row, seg in zip(a_2d, segment_ids)])

# 定位需要置0的分段
valid_segments = (sumseg > 0) & (sumseg < 2)

# 生成每个元素是否需要置0的掩码
fill_mask = valid_segments[np.arange(rows)[:, None], segment_ids]

# 直接修改原数组
a_2d[fill_mask] = 0

优化说明 #

• 用np.diff替代np.roll：np.roll会复制整个数组，np.diff仅计算相邻差异，内存占用更低、速度更快。
• 用np.add.reduceat替代列表推导求和：这是NumPy原生的分段求和函数，效率远高于Python循环。
• 二维数组的完全向量化版本彻底规避Python层循环，适合处理大规模数据集。

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