NumPy random.choice()函数的使用小结

np.random.choice() 是 NumPy 中用于随机抽样的核心函数，功能强大且灵活。以下是对其用法的小结：

一、基本语法

np.random.choice(a, size=None, replace=True, p=None)

参数说明：

• a：抽样的源
  • 整数 n：从 np.arange(n) 中抽样
  • 一维数组：直接从数组元素中抽样
• size：输出形状
  • None：返回单个值
  • 整数：返回一维数组
  • 元组：返回多维数组
• replace：是否放回
  • True：有放回抽样（默认）
  • False：无放回抽样
• p：概率分布
  • 与 a 等长的一维数组
  • 每个元素被选中的概率，需满足 sum(p)=1

二、基础用法示例

1. 从整数范围抽样

import numpy as np

# 从 0-4 中随机选一个数
np.random.choice(5)                     # 可能输出：3

# 从 0-9 中选3个（有放回）
np.random.choice(10, size=3)            # 可能输出：[7, 2, 7]

# 从 0-9 中选3个（无放回）
np.random.choice(10, size=3, replace=False)  # 可能输出：[9, 0, 4]

2. 从数组抽样

arr = ['a', 'b', 'c', 'd', 'e']

# 随机选择一个元素
np.random.choice(arr)                  # 可能输出：'c'

# 选择多个（可重复）
np.random.choice(arr, size=5)          # 可能输出：['a','e','d','a','c']

# 选择多个（不重复）
np.random.choice(arr, size=3, replace=False)  # 可能输出：['b','d','e']

三、高级用法

1. 指定概率分布

# 按指定概率抽样
np.random.choice(['A', 'B', 'C'], 
                 size=10, 
                 p=[0.1, 0.3, 0.6])  # A:10%, B:30%, C:60%
# 可能输出：['C','B','C','C','B','C','A','C','B','C']

2. 多维输出

# 生成 3x2 的随机矩阵
np.random.choice(10, size=(3, 2))  
# 可能输出：[[5, 1],
#           [9, 3],
#           [7, 0]]

3. 无放回抽样验证

# 无放回时，size 不能超过 a 的长度
np.random.choice(5, size=5, replace=False)  # 正确：随机排列
# 可能输出：[4, 0, 2, 1, 3]

# np.random.choice(5, size=6, replace=False)  # 错误！会报错

四、注意事项

概率归一化：p 参数会自动归一化，但建议自行保证 sum(p)=1

# 以下两种方式等价
np.random.choice([1, 2, 3], p=[0.2, 0.3, 0.5])
np.random.choice([1, 2, 3], p=[2, 3, 5])  # 自动归一化为 0.2,0.3,0.5

随机种子：为结果可重现，可设置随机种子

np.random.seed(42)
np.random.choice(10, size=3)  # 总是得到 [6, 3, 7]

新版随机数 API（推荐）：

# NumPy 1.17+ 推荐使用新 API
rng = np.random.default_rng(seed=42)
rng.choice(5, size=3)  # 更快的随机数生成器

五、实际应用场景

1. 模拟掷骰子

# 投掷 1000 次公平骰子
rolls = np.random.choice([1,2,3,4,5,6], size=1000)

# 投掷 1000 次不公平骰子（6点概率50%）
biased_rolls = np.random.choice([1,2,3,4,5,6], 
                                size=1000, 
                                p=[0.1,0.1,0.1,0.1,0.1,0.5])

2. 自助法（Bootstrap）抽样

data = np.array([1.2, 2.5, 3.1, 4.8, 5.3])
bootstrap_sample = np.random.choice(data, size=len(data), replace=True)

3. 随机划分数据集

data = np.arange(100)
train_idx = np.random.choice(100, size=80, replace=False)  # 80个训练索引
test_idx = np.setdiff1d(np.arange(100), train_idx)          # 剩余20个测试

六、与 random.choices() 对比（Python标准库）

# Python 标准库 equivalent
import random
random.choices(['a','b','c'], k=5, weights=[0.1,0.3,0.6])

掌握 np.random.choice() 可以有效替代许多需要循环实现的随机抽样场景，提高代码效率和可读性。