如何无循环高效地从数组中按条件随机迭代选取索引？

在开发人口统计模型时，需要一种高效方法从数组中按条件随机迭代选取索引。现有预期年龄分布数据：

import numpy as np
popsize = int(1e6)
age_pyramid_edges = np.linspace(0,100,101) # 0-100岁的单年龄段分组
age_pyramid_data = (np.array([0.01] * 100)*popsize).astype(int) # 假设每个年龄段占总人口1%

实际年龄数据：

actual_ages = np.random.uniform(0, 100, size=popsize)

需求是对比实际与预期年龄分布，当某年龄段A的实际人数超出预期M时，随机选取该年龄段的M个人员索引。目前采用循环实现如下，但速度较慢，希望找到无需循环的实现方式：

comparable_ages = np.digitize(actual_ages, age_pyramid_edges)-1 # 将年龄分段匹配预期分组
counts_of_actual_ages = np.bincount(comparable_ages, minlength=len(age_pyramid_edges)-1)
age_diffs = counts_of_actual_ages-age_pyramid_data

inds_to_flag = []
for age,age_diff in enumerate(age_diffs):
    if age_diff>0:
        inds_this_age = (comparable_ages==age).nonzero()[-1]
        inds = np.random.choice(inds_this_age, age_diff, replace=False).tolist()
        inds_to_flag.append(inds)

无循环高效实现方案 #

可以用numpy的向量化操作+随机排序替代循环，速度提升明显，具体思路：

1. 给所有样本打上所属年龄段的标签，算出每个年龄段需要筛选的样本数量
2. 将所有样本按年龄段分组，每组内部随机打乱顺序，直接截取每组超出预期数量的部分作为目标索引

代码实现如下：

import numpy as np

popsize = int(1e6)
age_pyramid_edges = np.linspace(0,100,101)
age_pyramid_data = (np.array([0.01] * 100)*popsize).astype(int)
actual_ages = np.random.uniform(0, 100, size=popsize)

# 1. 给年龄分段，得到每个样本的年龄段标签
comparable_ages = np.digitize(actual_ages, age_pyramid_edges) - 1

# 2. 生成所有样本索引，先按年龄段排序，再给每个年龄段内的索引随机洗牌
all_inds = np.arange(popsize)
# lexsort会先按第二个键（年龄段）排序，再按第一个键（随机序列）排序，实现组内随机打乱
shuffled_grouped_inds = all_inds[np.lexsort((np.random.permutation(popsize), comparable_ages))]

# 3. 计算每个年龄段的实际数量、要保留的数量，以及拆分索引用的累计位置
actual_counts = np.bincount(comparable_ages, minlength=100)
keep_counts = np.minimum(actual_counts, age_pyramid_data)
cum_actual = np.cumsum(actual_counts)

# 4. 拆分出每个年龄段的索引组，提取超出预期的部分
inds_to_flag = []
age_groups = np.split(shuffled_grouped_inds, cum_actual[:-1])
for keep_num, group in zip(keep_counts, age_groups):
    if len(group) > keep_num:
        inds_to_flag.append(group[keep_num:].tolist())

优化说明 #

• 避免了循环中反复生成布尔索引的开销，一次性完成所有索引的分组与随机打乱
• 全程使用numpy内置的向量化函数，相比Python循环效率提升显著，样本量越大优势越明显

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