如何用列表推导式与groupby实现Pandas DataFrame的子集筛选

核心思路 #

你的数据按patid分组存储（单患者对应多行记录），原isin方法会对每一行做匹配判断，而先通过groupby('patid')将同一患者的行聚合成组，再仅筛选符合patidList的组，能大幅减少重复判断次数——尤其当患者总数远少于总行数时，性能提升会非常明显。

实现代码 #

import pandas as pd

class TestClass():
    def __init__(self):
        dictExample = {
            "patid" : ["25001", "25001", "25001", "999", "999"],
            "interesting_stuff" : ["a","b","c","d","e"]
        }
        self.ioDF = pd.DataFrame(dictExample)

    def someMethod(self):
        patidList = ["25001"]
        # 转集合提速成员判断（集合判断复杂度O(1)，列表为O(n)）
        patid_set = set(patidList)
        # 列表推导式筛选目标组，再合并结果
        reducedDF = pd.concat([group for patid, group in self.ioDF.groupby('patid') if patid in patid_set])

testClass = TestClass()
testClass.someMethod()

性能优化补充 #

• 类型转换降本：如果patid是高重复率字符串，可提前转成category类型：self.ioDF['patid'] = self.ioDF['patid'].astype('category')，既能减少内存占用，也能加速groupby操作。
• 场景适配选择：
  • 当总行数多、患者数少（比如10万行对应10位患者），groupby方法优势显著；
  • 若患者数与行数接近，原isin方法可能更快（groupby本身有额外开销）。
• 极致性能尝试：还可以用query方法，语法简洁且部分场景下速度更优：

  reducedDF = self.ioDF.query('patid in @patidList')
  

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