熊猫数据框：当列总和超过阈值时统一缩小值

Question

初始情况

考虑以下示例数据框：

df = pd.DataFrame({
    'A': [3., 2., 1., np.nan],
    'B': [7., np.nan, 1., 3.],
    'C': [4., 5., 1., 2.],
    'D': [1., 0., 2., 3.]    
})

打印出来的样子是这样的：

     A    B    C    D
0  3.0  7.0  4.0  1.0
1  2.0  NaN  5.0  0.0
2  1.0  1.0  1.0  2.0
3  NaN  3.0  2.0  3.0

期望的结果

我现在想对该数据框的每一列执行以下操作：

1. 计算列值的总和（忽略任何 NaN 值）。
2. 如果总和超过 10.0，那么我想统一缩小列中的所有值，使新总和正好为 10.0（再次忽略任何 NaN 值）。

基本上我想获得一个如下所示的结果数据框：

     A         B         C    D
0  3.0  6.363636  3.333333  1.0
1  2.0       NaN  4.166667  0.0
2  1.0  0.909091  0.833333  2.0
3  NaN  2.727273  1.666667  3.0

到目前为止尝试过

下面的代码得到了想要的结果。

def helper_func(s):
    if s.sum() > 10.:
        return s * 10. / s.sum()
    else:
        return s

result_df = df.apply(helper_func)

但是，我觉得这段代码有点冗长且效率低下。根据我迄今为止对 pandas 的经验，我怀疑仍然可能有一个更加矢量化的解决方案。谁能帮我找到这个？

Answer 1

这是一种方法：

thres = 10
result = df * thres / df.sum().clip(lower=thres)

Answer 2

这是另一种方法：

colSums = df.sum()
df / ((colSums * (colSums > 10) / 10) + (colSums <= 10))

在这里，我们为每列创建一个具有总和值的变量，然后分母检查列的总和是否超过 10 调整这些列，使它们总和为 10。总和小于 10 的那些列将递增到 1，这样我们就不会除以 0。然后，生成的数组将跨列划分。这将返回所需的结果。

Out[46]: 
     A         B         C    D
0  3.0  6.363636  3.333333  1.0
1  2.0       NaN  4.166667  0.0
2  1.0  0.909091  0.833333  2.0
3  NaN  2.727273  1.666667  3.0