具有给定概率的切片/数组的Golang选项编号

时间:2018-05-24 10:53:14

标签: go probability

您好我正在寻找一种从给定概率向量的数组/切片中选择数字的方法,例如: 我们有数据[0,1,2]和概率向量[0.2,0.5,0.3] 所以我们选择0概率为0.2,1概率为0.5,2概率为0.3 在python我会使用numpy.random.choice。但我不知道如何在Go中做到这一点 我可以使用0-100之间的随机,然后使用if,如果数字为0-20,那么它的0和其他的相同。但我认为有更好的方法可以做到这一点,更通用的做它作为功能。

1 个答案:

答案 0 :(得分:1)

解决方案只是根据给定的概率(pdf)计算cdf,然后找到正确的存储桶:

func sample(cdf []float32) int {
    r := rand.Float32()

    bucket := 0
    for r > cdf[bucket] {
        bucket++
    }
    return bucket
}

通过

获取cdf
func main() {
    // probability density function
    pdf := []float32{0.3, 0.4, 0.2, 0.1}

    // get cdf
    cdf := []float32{0.0, 0.0, 0.0, 0.0}
    cdf[0] = pdf[0]
    for i := 1; i < 4; i++ {
        cdf[i] = cdf[i-1] + pdf[i]
    }

    // test sampling with 100 samples
    samples := []float32{0.0, 0.0, 0.0, 0.0}

    for i := 0; i < 100; i++ {
        samples[sample(cdf)]++
    }

    // normalize
    for i := 0; i < 4; i++ {
        samples[i] /= 100.
    }

    fmt.Println(samples)
    fmt.Println(pdf)

}

完整示例: https://play.golang.org/p/ceZx1EiYW3l

sample的输出遵循给定的概率:

samples: [0.33 0.43 0.15 0.09]
pdf: [0.3 0.4 0.2 0.1]