Pandas中的条件对平均值

Question

考虑以下表格。除days_old和SKU被视为整数外，所有其他列均视为布尔值。

pd.DataFrame({
    'SKU': [10,11,12,13,14,15],
    'frozen':[1,0,1,1,1,0],
    'vegetable':[1,1,0,1,1,0],
    'microwaveable':[1,0,0,0,0,1],
    'days_old':[21,9,11,2,6,14],
})

我需要一个共现计数表，该表能同时匹配列和索引变量，并给出days_old的平均值。期望的输出结果是：

AVERAGE DAYS OLD:
          frozen | vegetable | microwaveable 
frozen        10 |       9.6 |           21
vegetable    9.6 |       9.5 |           21
microw...     21 |        21 |         17.5

Jonathan Leffler · Answer

我理解你是想要求取这样的数据：当两列的值都为1时，这两列所对应的days_old值的均值。

为了实现这一目标，可以先对数据框进行调整，使得原来为1的值乘以对应的days_old值，而原来的0则替换为NaN。之后，利用Pandas的corr函数，该函数能高效地计算列间成对的相关性。由于corr允许使用自定义函数，并且它只处理两列中非NaN值的行对，我们可以简单地计算任一列的均值作为相关性的度量。需要注意的是，corr计算自相关时会在对角线上默认返回1，因此我们还需使用Numpy的fill_diagonal函数，将对角线上的值替换为各列的平均days_old值。

简而言之，就是先通过操作使得符合条件（两列值均为1）的数据点参与到平均值计算中，然后利用corr函数特性巧妙地完成任务，最后修正对角线数据以得到正确的列均值。

tmp = (df.drop(columns=['SKU', 'days_old'])
         .mul(df['days_old'], axis=0)
         .replace(0, np.nan)
      )

# get rid of rows with NaN, compute mean
out = tmp.corr(lambda a,b: a.mean())

# replace diagonal of 1s by the mean of the columns
np.fill_diagonal(out.values, tmp.mean())

较不高效的备选方案，采用 itertools.product：

from itertools import product

cols = df.columns.difference(['SKU', 'days_old'])

tmp = (df[cols]
         .mul(df['days_old'], axis=0)
         .replace(0, np.nan)
      )

out = (pd.Series({(a, b): tmp[a].where(tmp[b].notna()).mean()
                 for a,b in product(cols, repeat=2)})
         .unstack().reindex(index=cols, columns=cols)
      )

输出结果是：

               frozen  vegetable  microwaveable
frozen         10.000000   9.666667           21.0
vegetable       9.666667   9.500000           21.0
microwaveable  21.000000  21.000000           17.5

中间的临时变量 tmp 是：

   frozen  microwaveable  vegetable
0    21.0           21.0       21.0
1     NaN            NaN        9.0
2    11.0            NaN        NaN
3     2.0            NaN        2.0
4     6.0            NaN        6.0
5     NaN           14.0        NaN

Sarfaraz Nawaz · Answer

另一种方法，不先对数据框进行乘法操作：

from itertools import product

out = pd.Series(
    {
        (a, b): df.loc[df[a].astype(bool) & df[b].astype(bool), "days_old"].mean()
        for a, b in product(df.columns[1:-1], df.columns[1:-1])
    },
).unstack(1)

print(out)

输出结果为：

                  frozen  microwaveable  vegetable
frozen         10.000000           21.0   9.666667
microwaveable  21.000000           17.5  21.000000
vegetable       9.666667           21.0   9.500000

Matt Johnson-Pint · Answer

一种方法是使用点乘运算（dot product）来实现：

# 选取相关的列
df2 = df[['frozen', 'vegetable', 'microwaveable']]

# 计算列与'days_old'的乘积，然后转置并点乘自身，再除以自身转置后点乘自身的结果，最后四舍五入到两位小数
result = df2.T.dot(df2.mul(df['days_old'], axis=0)).div(df2.T.dot(df2)).round(2)

输出结果为：

               frozen  vegetable  microwaveable
frozen          10.00       9.67           21.0
vegetable        9.67       9.50           21.0
microwaveable   21.00      21.00           17.5