Pandas基于多列条件筛选行并统计匹配数的更优实现方法

性能优化方案 #

单次查询场景（仅执行少数几次匹配） #

• 优先替换原生Python sum()为pandas内置矢量化求和方法，避免Python层遍历元素，性能可提升3~5倍：

num_occurrence = ((df["Chr"] == chrom) & 
                  (df["Start"] == int(position)) & 
                  (df["Alt"] == allele)).sum()

• 追求更简洁的写法可以用query()方法，可读性更高，大样本下性能和上述写法持平甚至更优：

# @ 符号用于引用当前作用域的外部变量
num_occurrence = df.query("Chr == @chrom and Start == @position and Alt == @allele").shape[0]

多次查询场景（需要反复执行多条件匹配） #

如果需要频繁做这类统计，提前对查询字段预处理可以把查询速度提升1~2个数量级：

1. 方案1：建复合索引，查询时走索引匹配，时间复杂度从O(n)降到O(logn)

# 首次运行时先建索引，只需执行一次
df = df.set_index(["Chr", "Start", "Alt"])

# 后续查询逻辑
try:
    num_occurrence = df.loc[(chrom, int(position), allele)].shape[0]
except KeyError:
    # 无匹配结果时返回0
    num_occurrence = 0

2. 方案2：预计算所有组合的出现次数，查询时直接O(1)取值，性能最优

# 首次运行预计算所有三字段组合的计数，只需执行一次
count_map = df.value_counts(["Chr", "Start", "Alt"])

# 后续查询直接取值即可，无匹配时默认返回0
num_occurrence = count_map.get((chrom, int(position), allele), 0)

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