Pandas Lookup 即将被弃用，如何优雅地替代它？

2024-07-14 05:30:25

如何优雅地根据条件选择列值填充新列？

在数据处理过程中，我们常常需要根据特定条件从其他列中提取值来填充新列。Pandas 的 lookup 函数曾是解决这类问题的一把利器，但很遗憾，它即将被弃用。官方推荐使用 .melt 和 .loc 方法进行替代，但这真的足够优雅和高效吗？

本文将深入探讨如何以更优雅、更高效的方式替代 lookup 函数，并结合实际代码示例，帮助你轻松应对类似问题。

从一个例子开始

假设我们有一个 DataFrame，其中包含两列："B" 列存储类别标签，"group" 列代表分组信息。我们的目标是根据每行的 "group" 值和 "B" 列中的类别标签，从一个预先计算好的 DataFrame 中查找对应值，并填充到新列 "count" 中。

import pandas as pd

df = pd.DataFrame({'B': ['X', 'X' , 'Y', 'X', 'Y', 'Y', 
                         'X', 'X', 'Y', 'Y', 'X', 'Y'],
                   'group': ["IT", "IT", "IT", "MV", "MV", "MV", 
                             "IT", "MV", "MV", "IT", "IT", "MV"]})

# 预先计算的 DataFrame
a = (pd.concat([df, df['B'].str.get_dummies()], axis=1)
     .groupby('group').rolling(3, min_periods=1).sum()
     .sort_index(level=1).reset_index(drop=True))  

# 使用 lookup 函数 (即将被弃用)
df['count'] = a.lookup(df.index, df['B'])

上述代码使用 lookup 函数简洁地实现了目标，但会触发 FutureWarning，提示我们该方法即将被弃用。官方建议的替代方案如下：

b = pd.melt(a, value_vars=a.columns, var_name='B', ignore_index=False)
b.index.name='index'
df.index.name='index'
df = df.merge(b, on=['index','B'])

诚然，这种方法可以实现目标，但代码显得冗长，可读性较差，效率也不够理想。

另辟蹊径：利用索引对齐

Pandas 的强大之处在于其灵活的索引机制，我们可以利用索引对齐和多级索引功能，以更优雅、更高效的方式解决这个问题。

多级索引：化解查找难题

让我们先将预先计算的 DataFrame "a" 设置为多级索引，以 "group" 和 "B" 列作为索引。

a = a.set_index(['group', 'B'], drop=True)

现在，DataFrame "a" 的索引结构可以理解为一个二维表格，"group" 和 "B" 列的值组合构成了唯一的行索引。

`loc` 方法：精准定位目标值

利用设置好的多级索引，我们可以使用 loc 方法直接从 DataFrame "a" 中查找对应值。

df['count'] = a.loc[zip(df['group'], df['B'])].values

zip(df['group'], df['B']) 将 "group" 和 "B" 列的值打包成元组列表，每个元组对应 DataFrame "a" 中的一个唯一行索引。loc 方法根据这些索引从 "a" 中提取对应值，最终赋值给 DataFrame "df" 的新列 "count"。

完整代码及解析

import pandas as pd

df = pd.DataFrame({'B': ['X', 'X' , 'Y', 'X', 'Y', 'Y', 
                         'X', 'X', 'Y', 'Y', 'X', 'Y'],
                   'group': ["IT", "IT", "IT", "MV", "MV", "MV", 
                             "IT", "MV", "MV", "IT", "IT", "MV"]})

a = (pd.concat([df, df['B'].str.get_dummies()], axis=1)
     .groupby('group').rolling(3, min_periods=1).sum()
     .sort_index(level=1).reset_index(drop=True))        

# 设置多级索引
a = a.set_index(['group', 'B'], drop=True)

# 使用 loc 方法和多级索引查找
df['count'] = a.loc[zip(df['group'], df['B'])].values

print(df)