获取pandas groupby transform中的组名

Question

我想做的是这样的。我在pandas中有一个数据框（DataFrame）：

import numpy as np
import pandas as pd

n_cols = 3
n_samples = 4
df = pd.DataFrame(np.arange(n_samples * n_cols).reshape(n_samples, n_cols), columns=list('ABC'))
print(df)

输出结果：

   A   B   C
0  0   1   2
1  3   4   5
2  6   7   8
3  9  10  11

每一行数据都有一个类别：

cat = pd.Series([1,1,2,2])

而且我有一个与每个类别相关的参考行：

df_ref = pd.DataFrame(np.zeros((2, n_cols)), index=[1,2], columns=list('ABC'))
df_ref.loc[1] = 10
print(df_ref)

输出结果：

      A     B     C
1  10.0  10.0  10.0
2   0.0   0.0   0.0

我想用更优雅的方式来完成以下操作（比如，使用groupby和transform）：

result = df.copy()
for i in range(n_cols):
  result.iloc[i] = df.iloc[i] - df_ref.loc[cat[i]]
print(results)

输出结果：

    A   B   C
0 -10  -9  -8
1  -7  -6  -5
2   6   7   8
3   9  10  11

我觉得像这样应该可以工作：

df.groupby(cat).transform(lambda x: x - df_ref.loc[x.GROUP_NAME])

这里的x.GROUP_NAME是用来获取正在进行transform操作的组的名称。在pandas的transform文档中写道：“每个组都有一个‘name’属性，以便你知道自己在处理哪个组。”我尝试访问x.name，但那给出的却是列的名称，而不是组的名称。所以我不明白这个文档所指的是什么。

Answer 1

如果我理解正确的话，你可以这样做：

df = df.set_index(cat)
df = df - df_ref

print(df)

输出结果是：

      A     B     C
1 -10.0  -9.0  -8.0
1  -7.0  -6.0  -5.0
2   6.0   7.0   8.0
2   9.0  10.0  11.0

Answer 2

其实不需要用到 groupby，只要用 reindex 来调整 df_ref 的索引，然后转换成数组就可以了：

df -= df_ref.reindex(cat).values

或者，如果你想要一个副本的话：

out = df.sub(df_ref.reindex(cat).values)

需要注意的是，你的方法也可以用 groupby.apply 来实现：

out = df.groupby(cat, group_keys=False).apply(lambda x: x - df_ref.loc[x.name])

输出结果：

      A     B     C
0 -10.0  -9.0  -8.0
1  -7.0  -6.0  -5.0
2   6.0   7.0   8.0
3   9.0  10.0  11.0