Python中含层级字符串的for循环向量化实现方案咨询

需求说明 #

统计start与stop组成的区间内apple和orange的数量，额外规则：当第2个start-stop区间结束时，将当前的apple计数翻倍。原代码用循环实现，现在改用numpy向量化方式优化。

原代码逻辑回顾 #

• 仅当处于start和stop之间时，才统计apple和orange
• 每完成一个start-stop区间，计数器加1，第2个区间完成时，apple计数乘以2
• 原代码运行结果：count_apple=3，count_orange=2

向量化实现步骤 #

1. 标记关键元素位置

把原始数组转为numpy数组，生成布尔数组标记各类关键元素的位置：

import numpy as np

arr = np.array(["start","apple","orange","stop","chery","start","chery","orange","stop","apple","start","chery","apple","stop","apple"])

# 标记各类元素的布尔数组
is_start = arr == "start"
is_stop = arr == "stop"
is_apple = arr == "apple"
is_orange = arr == "orange"

2. 生成有效区间掩码

通过累积和确定每个元素所属的分组编号，再用反向累积和筛选出start到stop之间的有效元素：

# 计算每个元素的分组编号（每次遇到start，分组号+1）
group_id = np.cumsum(is_start)
# 反向统计stop的累积和，标记stop之后的元素为无效
reverse_stop_cumsum = np.cumsum(is_stop[::-1])[::-1]
# 有效元素需满足：属于某个分组（group_id>0）且未到达对应分组的stop
mask = (group_id > 0) & (reverse_stop_cumsum == 0)

3. 按分组统计元素数量

用np.bincount统计每个分组内的apple和orange数量：

# 统计每个分组的apple数量
apple_per_group = np.bincount(group_id[mask & is_apple], minlength=np.max(group_id)+1)
# 统计每个分组的orange数量
orange_per_group = np.bincount(group_id[mask & is_orange], minlength=np.max(group_id)+1)

4. 应用第2组翻倍规则

计算总计数，并处理第2组的apple翻倍逻辑：

count_apple = np.sum(apple_per_group)
# 若存在第2组，将该组的apple数量额外加一次（等价于原计数乘以2）
if len(apple_per_group) >= 2:
    count_apple += apple_per_group[1]

count_orange = np.sum(orange_per_group)

print(f"count_apple: {count_apple}, count_orange: {count_orange}")
# 输出：count_apple: 3, count_orange: 2

完整向量化代码 #

import numpy as np

arr = np.array(["start","apple","orange","stop","chery","start","chery","orange","stop","apple","start","chery","apple","stop","apple"])

is_start = arr == "start"
is_stop = arr == "stop"
is_apple = arr == "apple"
is_orange = arr == "orange"

group_id = np.cumsum(is_start)
reverse_stop_cumsum = np.cumsum(is_stop[::-1])[::-1]
mask = (group_id > 0) & (reverse_stop_cumsum == 0)

apple_per_group = np.bincount(group_id[mask & is_apple], minlength=np.max(group_id)+1)
orange_per_group = np.bincount(group_id[mask & is_orange], minlength=np.max(group_id)+1)

count_apple = np.sum(apple_per_group)
if len(apple_per_group) >= 2:
    count_apple += apple_per_group[1]

count_orange = np.sum(orange_per_group)

print(count_apple, count_orange)

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