如何使用pandas按概率权重为不同国家用户随机分配对应城市？

核心是先将表1按国家聚合得到每个国家对应的城市和权重映射，再根据用户所属国家按权重随机采样城市。代码兼容多国家、任意城市数量的场景。

完整代码示例 #

首先导入依赖：

import pandas as pd
import numpy as np

步骤1：准备数据 #

你可以替换下方的示例数据为自己的真实数据读取逻辑（比如用pd.read_csv读入本地文件）：

# 构造表1（国家-城市概率表）示例数据
df1 = pd.DataFrame({
    'country': ['美国', '美国', '美国', '中国', '中国'],
    'city': ['纽约', '休斯顿', '洛杉矶', '北京', '上海'],
    'probability': [0.3, 0.5, 0.2, 0.6, 0.4]
})

# 构造表2（用户表）示例数据
df2 = pd.DataFrame({
    'user_id': [1,2,3,4,5,6,7,8,9,10],
    '所属国家': ['美国', '美国', '美国', '美国', '美国', '中国', '中国', '中国', '中国', '中国']
})

步骤2：生成国家-城市权重映射 #

# 按国家分组，存储每个国家对应的城市列表、概率列表
country_city_map = df1.groupby('country').apply(
    lambda x: (x['city'].tolist(), x['probability'].tolist())
).to_dict()

步骤3：为用户分配城市 #

小数据量实现（万级及以下用户，写法简单）

逐行遍历用户分配城市：

df2['city'] = df2.apply(
    lambda row: np.random.choice(
        country_city_map[row['所属国家']][0],
        p=country_city_map[row['所属国家']][1]
    ), axis=1
)

大数据量优化实现（十万级及以上用户，性能提升明显）

按国家批量采样，避免逐行遍历的性能损耗：

df2 = df2.groupby('所属国家', group_keys=False).apply(
    lambda g: g.assign(
        city = np.random.choice(
            country_city_map[g.name][0],
            size = len(g),
            p = country_city_map[g.name][1]
        )
    )
)

注意事项 #

• 请确保表1中每个国家的probability求和为1，否则np.random.choice会抛出异常。如果原始概率未归一化，可以在生成映射前先做归一化处理：

# 概率归一化处理示例
df1['probability'] = df1.groupby('country')['probability'].apply(lambda x: x / x.sum())

• 如果需要固定每次运行的分配结果，可以在运行代码前设置随机种子：np.random.seed(42)（数字可自行修改）

内容的提问来源于stack exchange，提问作者Callum Matthews