使用 Pandas .apply() 创建多列

Question

我有两个 pandas DataFrame，它们都包含相同的类别但不同的“id”列。为了说明，第一个表如下所示：

df = pd.DataFrame({
    'id': list(np.arange(1, 12)),
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'c', 'c', 'c', 'c', 'c'],
    'weight': list(np.random.randint(1, 5, 11))
})

df['weight_sum'] = df.groupby('category')['weight'].transform('sum')
df['p'] = df['weight'] / df['weight_sum']

Output:

id  category    weight  weight_sum  p
0   1   a   4   14  0.285714
1   2   a   4   14  0.285714
2   3   a   2   14  0.142857
3   4   a   4   14  0.285714
4   5   b   4   8   0.500000
5   6   b   4   8   0.500000
6   7   c   3   15  0.200000
7   8   c   4   15  0.266667
8   9   c   2   15  0.133333
9   10  c   4   15  0.266667
10  11  c   2   15  0.133333

第二个只包含'id'和'category'。

我要做的是创建一个第三个数据帧，它将继承第二个数据帧的id，加上第一个数据帧的ids 的三个新列- 应根据p 列选择每个，这表示其在该类别中的权重。

我尝试了多种解决方案，并正在考虑应用 np.random.choice 和 .apply()，但无法找到实现该功能的方法。

编辑：

所需的输出类似于：

user_id id_1    id_2    id_3
0   2   3   1   2
1   3   2   2   3
2   4   1   3   1

每个id 是根据其概率和各自的category 选择的（两个DataFrame 都有此列），并且相同的user_id 不会出现多次。

Desired DataFrame

Answer 1

IIUC，您想选择具有加权概率的同一类别的随机 ID。为此，您可以构建一个辅助数据框 (dfg) 并使用 apply:

df2 = pd.DataFrame({
    'id': np.random.randint(1, 12, size=11),
    'category': ['a', 'a', 'a', 'a', 'b', 'b', 'c', 'c', 'c', 'c', 'c']})

dfg = df.groupby('category').agg(list)

df3 = df2.join(df2['category']
               .apply(lambda r: pd.Series(np.random.choice(dfg.loc[r, 'id'],
                                                           p=dfg.loc[r, 'p'],
                                                           size=3)))
               .add_prefix('id_')
               )

输出：

    id category  id_0  id_1  id_2
0   11        a     2     3     3
1   10        a     2     3     1
2    4        a     1     2     3
3    7        a     2     1     4
4    5        b     6     5     5
5   10        b     6     5     6
6    8        c     9     8     8
7   11        c     7     8     7
8   11        c    10     8     8
9    4        c     9    10    10
10   1        c    11    11     9