python - 迭代 pandas df rows 以根据条件计算新变量

Question

我有一个如下所示的数据框:

dict1 = {'category': {0: 0.0, 1: 1.0, 2: 0.0, 3: 1.0, 4: 0.0},
 'Id': {0: 24108, 1: 24307, 2: 24307, 3: 24411, 4: 24411},
 'count': {0: 3, 1: 2, 2: 33, 3: 98, 4: 33},
 'weight': {0: 0.5, 1: 0.2, 2: 0.7, 3: 1.2, 4: 0.39}}

df1 = pd.DataFrame(dict1)

    category    Id  count   weight
0   0.0      24108    3     0.50
1   1.0      24307    2     0.20
2   0.0      24307   33     0.70
3   1.0      24411   98     1.20
4   0.0      24411   33     0.39

大约有 1000 个这样的条目。有些只有 1 个 category 标签，例如 Id 24108。其他 Id 有两个条目，因为有两个 category 标签，例如 Id 24307 和 24411。

我想生成一个名为 val 的新列，它遵循 3 条规则之一，具体取决于某些条件。

如果 Id 只有 1 个与其关联的标签，例如 Id 24108，则新列 val 中的值应为计数和重量之间的简单乘积。

如果 Id 有 2 个标签，例如 Id 24307 或 24411，则脚本应查看两个条目中哪一个具有较大的 计数 值(value)第一。

对于 Id 24307，category 0 的计数高于 category 1 的计数，因此 val code> 列应为 category 1 旁边的 2 * (0.2+1) 和 category 0

旁边的 33 * (1-0.7)

对于 Id 24411，category 0 的计数低于 category 1 的计数，因此 val code> 列应为 category 1 旁边的 98 * (1-1.2) 和 category 0 旁边的 33 * (0.3+1)。

预期的输出应该是这样的:

    category    Id  count   weight  val
0   0.0      24108    3     0.50   1.50
1   1.0      24307    2     0.20   2.40
2   0.0      24307   33     0.70   9.90
3   1.0      24411   98     1.20 -19.60
4   0.0      24411   33     0.39  42.90

Answer 1

您可以使用 bool 算术和numpy.select:

g = df1.groupby('Id')

# number of categories
n = g['category'].transform('nunique')
# max count
m = g['count'].transform('max')

# selection based on conditions
df1['val'] = np.select(
          [n.eq(1), df1['count'].eq(m)], # case if unique, case if max
          # formula for unique, formula for max
          [df1['count']*df1['weight'], df1['count']*(1-df1['weight'])],
          # default value (i.e. not unique and non-max count)
          df1['count']*(1+df1['weight'])
         )

输出:

   category     Id  count  weight    val
0       0.0  24108      3    0.50   1.50
1       1.0  24307      2    0.20   2.40
2       0.0  24307     33    0.70   9.90
3       1.0  24411     98    1.20 -19.60
4       0.0  24411     33    0.39  45.87

如果您确实只有 2 个类别，那么您可以简化算术，将非最大值视为最小值。它有点很老套，但应该可以正常工作:

g = df1.groupby('Id')

# 0 if unique category else 1
n = g['category'].transform('nunique').gt(1).astype(int)
# -1 if count is min of group else 1
m = (df1['count'] != g['count'].transform('min'))*2-1

df1['val'] = df1['count'] * (n-m*df1['weight'])