Python中从多维联合概率质量函数采样索引的实现方法

先回顾一维场景的实现：给定非负归一化向量p，按p[k]的概率采样索引k，可以用np.random.choice实现：

import numpy as np

p = [0.2, 0.3, 0.1, 0.3, 0.1]
indices = np.arange(len(p))
k = np.random.choice(indices, p=p)

三维张量的采样方案 #

对于三维非负归一化张量p（形状为I×J×K），要按p[i,j,k]的概率采样索引(i,j,k)，核心思路是将多维张量展平为一维，利用一维采样逻辑完成后，再将一维索引还原为多维坐标：

import numpy as np

# 构造三维归一化张量示例
I, J, K = 2, 3, 4
p = np.random.rand(I, J, K)
p /= p.sum()  # 确保张量整体归一化

# 1. 将张量展平为一维
p_flat = p.flatten()
# 2. 生成一维索引范围
flat_indices = np.arange(p_flat.size)
# 3. 采样一维索引
sampled_flat_idx = np.random.choice(flat_indices, p=p_flat)
# 4. 将一维索引转换为三维坐标(i,j,k)
sampled_idx = np.unravel_index(sampled_flat_idx, p.shape)

print(f"采样得到的三维索引：{sampled_idx}")
print(f"对应概率值：{p[sampled_idx]}")

关键说明 #

• np.flatten()将多维张量转换为一维数组，保留了原张量元素的顺序映射
• np.unravel_index()是np.flatten()的逆操作，能根据原张量的形状，把一维索引还原为对应的多维坐标
• 必须确保整个三维张量的元素和为1（即全局归一化），否则采样概率会不符合预期

内容的提问来源于stack exchange，提问作者Sakurai.JJ