在for循环中修改数据框（DataFrame）

Question

我在审查一段Python代码（由于不常使用Python，所以不确定设计模式中的细微差别），看到以下的注释：

# Mutating dataframes in a loop isn't efficient and stylistic. TODO: find alternative

这段注释针对的是如下代码：

for id in ids:
  rows4id = df.id == id
  date2scale = get_date2scale(date4id[id])
  df.loc[rows4id, "scale"] = df.loc[rows4id, "date"].map(date2scale)

如何重写这段代码以使其“更高效且风格更加优雅”？

示例数据：

import pandas as pd

df = pd.DataFrame({
  "id": [0, 1, 2],
  "date": [10, 10, 10],
  "scale": [0, 0, 0]
})

ids = [1, 2]

def get_date2scale(id):
  if id == 1:
    return {10: 5.1}
  elif id == 2:
    return {10: 5.2}
  else:
    raise ValueError("Should not be called for this id")

期望的结果是：

df = pd.DataFrame({
  "id": [0, 1, 2],
  "date": [10, 10, 10],
  "scale": [0, 5.1, 5.2]
})

Vlad Bezden · Answer

这段代码本质上执行的是：

# 针对每个id
for id in ids:  
  # 获取匹配行的掩码
  rows4id = df.id == id 
  # 获取当前id对应的映射器
  date2scale = get_date2scale(date4id[id])
  # 将所有匹配行替换为映射后的值
  df.loc[rows4id, "scale"] = df.loc[rows4id, "date"].map(date2scale)

可以将其替换为先使用loc进行筛选，然后应用groupby.apply方法：

m = df['id'].isin(ids)

df.loc[m, 'scale'] = (df[m].groupby('id', group_keys=False)['date']
                           .apply(lambda g: g.map(get_date2scale(g.name)))
                     )

然而，这种方法在效率上可能并不会显著提升，因为它实质上执行了相同的循环操作。除非ids中包含数据框中不存在的id（在这种情况下，它应该会更快）。

ecatmur · Answer

使用mask方法：

idx = [1, 2]
cond = df['id'].isin(idx) & (df['date'] == 10)
df['scale'] = df['scale'].mask(cond, df['id'] * 0.1 + 5)

# 输出结果：
#    id  date  scale
# 0   0   10     0.0
# 1   1   10     5.1
# 2   2   10     5.2

这样无需定义函数和循环。

如果情况更为复杂，且上述示例中的函数仅作为占位符，请考虑使用numpy.select：

import numpy as np

cond1 = df['id'] == 1
cond2 = df['id'] == 2
df['scale'] = np.select([cond1, cond2], [df['date'].replace({10: 5.1}), df['date'].replace({10: 5.2})], df['scale'])

# 输出结果与上述相同