如何将itertools的all_equal函数应用于Pandas Series生成布尔列

问题背景 #

我是一名外科医生，拥有患者手术数据集mo，其中OPERTN_01_list列存储患者的手术代码列表（共4种代码），其value_counts统计结果如下：

['O011']                                    2785
['O012']                                    1813
['O011', 'O011']                             811
['O013']                                     532
['O012', 'O012']                             522
['O014']                                     131
['O013', 'O013']                             125
['O014', 'O014']                              26
['O012', 'O011']                              24
['O011', 'O012']                              20
['O011', 'O011', 'O011']                      14
['O011', 'O013']                              12
['O012', 'O012', 'O011']                       6
['O011', 'O012', 'O012']                       6
['O011', 'O011', 'O011', 'O011']               5
['O013', 'O013', 'O013']                       5
['O013', 'O011']                               4
['O012', 'O012', 'O012']                       4
['O012', 'O013']                               3
['O013', 'O014']                               3
['O011', 'O013', 'O013']                       3
['O012', 'O014']                               3
['O011', 'O012', 'O011']                       2
['O012', 'O013', 'O013']                       2
['O011', 'O014']                               2
['O013', 'O012', 'O012']                       2
['O014', 'O014', 'O014']                       2
['O013', 'O012']                               1
['O012', 'O012', 'O013', 'O013', 'O013']       1
['O012', 'O011', 'O012']                       1
['O011', 'O011', 'O012']                       1
['O013', 'O013', 'O011']                       1
['O011', 'O011', 'O012', 'O012']               1
['O014', 'O013', 'O013']                       1
['O013', 'O013', 'O012']                       1
['O012', 'O011', 'O011']                       1
['O011', 'O012', 'O013']                       1
['O013', 'O011', 'O011']                       1
['O012', 'O012', 'O012', 'O012']               1
['O013', 'O013', 'O012', 'O012']               1
['O014', 'O013', 'O011', 'O011']               1
['O012', 'O011', 'O011', 'O011']               1
['O013', 'O011', 'O012']                       1

我需要新增一个布尔列，判断每位患者的所有手术代码是否完全相同。参考itertools方案编写了如下函数：

from itertools import groupby

def all_equal(iterable):
    g = groupby(iterable)
    return next(g, True) and not next(g, False)

但直接执行mo['eq'] = all_equal(mo['OPERTN_01_list'])后，新列mo['eq']所有值均为False，无法实现需求。

错误原因 #

直接将整个Series传入all_equal函数时，函数会迭代Series中的每个列表元素（即把每个列表当作iterable的一项），而非对每个单独的列表进行元素一致性判断，导致结果不符合预期。

解决方案 #

方法1：使用apply方法调用现有函数 #

对OPERTN_01_list列的每个元素单独应用all_equal函数：

mo['eq'] = mo['OPERTN_01_list'].apply(all_equal)

方法2：更直观的自定义函数 #

如果觉得groupby的写法不够直观，可以改用以下函数，逻辑更清晰：

def all_equal(lst):
    # 处理空列表的情况（如果存在）
    if not lst:
        return True
    # 判断所有元素是否等于第一个元素
    return all(x == lst[0] for x in lst)

mo['eq'] = mo['OPERTN_01_list'].apply(all_equal)

方法3：利用集合特性简化代码 #

利用集合元素唯一性的特点，一行代码实现需求：

# 空列表返回True，否则判断集合长度是否为1
mo['eq'] = mo['OPERTN_01_list'].apply(lambda x: len(set(x)) <= 1)

以上三种方法都能正确判断每个列表中的手术代码是否完全相同，比如['O011']、['O011','O011']会返回True，['O011','O012']会返回False，符合需求。

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