具有给定概率的切片/阵列中的Golang选择数

Hello i'm looking for a way to choose number from array/slice with given probability vector, for example like this: We have data [0,1,2] and probability vector[0.2,0.5,0.3] so we choose 0 with probability 0.2, 1 with probability 0.5 and 2 with probability 0.3 in python i would use numpy.random.choice . but i don't know how to do it in Go i can use random from 0-100 and then with if's make something like if number is 0-20 then its 0 and same way for others. But i think there is better way to do it and more generic to do it as function.

The solution is simply to compute the cdf from the given probabilities (pdf) and then find the correct bucket:

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

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

Getting the cdf is done by

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)

}

Full example: https://play.golang.org/p/ceZx1EiYW3l

The output of sample follows the given probabilities:

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