无需循环实现列表/NumPy数组的零值填充与区间计数

嘿，针对你处理大规模数组的需求，这里有两个完全基于NumPy矢量化操作的解决方案，全程不需要写for循环，效率拉满！

先把输入数据转换成NumPy数组（毕竟要高效处理，NumPy是首选）：

import numpy as np

ll = np.array([7.2, 0, 0, 0, 0, 0, 6.5, 0, 0, -8.1, 0, 0, 0, 0])
i = np.array([0, 6, 9])  # 非零元素的索引

生成填充数组a（用前一个非零元素填充零值） #

核心思路是计算每个非零元素需要重复的次数，然后用np.repeat一次性生成结果：

1. 先计算每个非零元素对应的区间长度：把最后一个元素的下一个位置（数组长度）加到索引列表末尾，然后用diff计算相邻索引的差，就是每个区间的元素个数。
2. 用np.repeat把每个非零元素重复对应次数，直接得到填充后的数组。

代码实现：

# 计算每个非零元素对应的区间长度
lengths = np.diff(np.append(i, len(ll)))
# 生成填充数组a
a = np.repeat(ll[i], lengths)

运行后a的结果就是：array([7.2, 7.2, 7.2, 7.2, 7.2, 7.2, 6.5, 6.5, 6.5, -8.1, -8.1, -8.1, -8.1, -8.1])，完全符合你的需求。

生成区间计数数组b（每个区间从0开始计数） #

这里用矢量化的累积和与索引映射来实现，同样没有循环：

1. 先创建一个标记数组，在非零元素的索引位置设为1，其他为0，然后计算累积和，得到每个元素所属的段ID。
2. 用每个元素的全局索引减去对应段的起始索引，就能得到该段内的计数。

代码实现：

# 生成段ID数组：每个非零位置触发段ID加1
segment_ids = np.zeros_like(ll, dtype=int)
segment_ids[i] = 1
segment_ids = np.cumsum(segment_ids)
# 生成每个段的起始索引数组，和原数组长度匹配
start_indices = np.repeat(i, lengths)
# 计算每个位置的区间计数
b = np.arange(len(ll)) - start_indices

运行后b的结果是：array([0, 1, 2, 3, 4, 5, 0, 1, 2, 0, 1, 2, 3, 4])，完美对应示例结果。

为什么这样高效？ #

NumPy的矢量化操作是在底层用C语言实现的，避开了Python解释器的循环开销，哪怕是处理百万级别的数组，速度也比Python的for循环快几十甚至上百倍，非常适合你的大规模数据场景。

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