基于列过滤器的pandas数据聚合

Question

我有一个这样的数据框

col1  col2  col3  col4  col5
A     A1    X     1      2
A     A2    Y     2      2
A     A3    Z     1      2
B     B1    X     2      2
B     B2    Y     2      2
B     B3    Z     1      2

我有这样的映射。

映射

col1  col3
 A --> X
 A --> Y
(A,B) --> Z
(A,B) --> (Y,Z)

这个映射就像我们在 excel 中使用的过滤器选项一样工作。 例如：对于 A-->X，我们将从 col1 中过滤 A，从 col3 中过滤 X，然后从 col4 中得到总和

对于 (A,B)-->Z，我们将从 col1 中过滤 A 和 B，从 col3 中过滤 Z，无论我们从 col4 中得到什么总和都应该被存储。

我尝试过 groupby 和 pivot_table 方法，但无法涵盖映射中的所有情况。

df1 = df.pivot_table(index= [col1,col3], columns = col3, values = col4, aggfunc='sum')

我的预期结果是。

col1  X   Y   Z   YZ
A     1   2   2   6
B             2   6

有什么方法可以使用 groupy 或 pivot_table 得到这个，请告诉我。

Answer 1

使用您提供的数据框：

import pandas as pd

df = pd.DataFrame(
    {
        "col1": ["A", "A", "A", "B", "B", "B"],
        "col2": ["A1", "A2", "A3", "B1", "B2", "B3"],
        "col3": ["X", "Y", "Z", "X", "Y", "Z"],
        "col4": [1, 2, 1, 2, 2, 1],
        "col5": [2, 2, 2, 2, 2, 2],
    }
)

这是使用 Pandas groupby 和 advanced indexing 的一种方法：

# Get values
df = df.groupby(["col1", "col3"]).agg(list)["col4"].apply(lambda x: x[0])

print(df)
# Output

col1  col3
A     X       1
      Y       2
      Z       1
B     X       2
      Y       2
      Z       1

# Assign values to new dataframe
new_df = pd.DataFrame()

for i in ("X", "Y"):
    new_df.at["A", i] = df.loc[("A", i)]

for i in ("A", "B"):
    new_df.at[i, "Z"] = df.loc["A":"B", "Z", :].reset_index().sum()["col4"]
    new_df.at[i, "YZ"] = df.loc["A":"B", "Y":"Z", :].reset_index().sum()["col4"]

# Cleanup
new_df = new_df.fillna(0).astype(int).replace(0, "")
new_df.columns.name = "col1"

然后：

print(new_df)

# Output
col1  X  Y  Z  YZ
A     1  2  2   6
B           2   6